Title: Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep Learning Theory

Feature Learning

Dynamics

Ensembles
Abstract: We analyze the dynamics of finite width effects in wide but finite feature learning neural networks. Starting from a dynamical mean field theory  description of infinite width deep neural network kernel and prediction dynamics, we provide a characterization of the $\mathcal{O}(1/\sqrt{\text{width}})$ fluctuations of the DMFT order parameters over random initializations of the network weights. Our results, while perturbative in width, unlike prior analyses, are non-perturbative in the strength of feature learning. In the lazy limit of network training, all kernels are random but static in time and the prediction variance has a universal form. However, in the rich, feature learning regime, the fluctuations of the kernels and predictions are dynamically coupled with a variance that can be computed self-consistently. In two layer networks, we show how feature learning can dynamically reduce the variance of the final tangent kernel and final network predictions. We also show how initialization variance can slow down online learning in wide but finite networks. In deeper networks, kernel variance can dramatically accumulate through subsequent layers at large feature learning strengths, but feature learning continues to improve the signal-to-noise ratio of the feature kernels. In discrete time, we demonstrate that large learning rate phenomena such as edge of stability effects can be well captured by infinite width dynamics and that initialization variance can decrease dynamically. For CNNs trained on CIFAR-10, we empirically find significant corrections to both the bias and variance of network dynamics due to finite width.

Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks

Abstract

We analyze the dynamics of finite width effects in wide but finite feature learning neural networks. Starting from a dynamical mean field theory description of infinite width deep neural network kernel and prediction dynamics, we provide a characterization of the O (1 / √ width ) fluctuations of the DMFT order parameters over random initializations of the network weights. Our results, while perturbative in width, unlike prior analyses, are non-perturbative in the strength of feature learning. In the lazy limit of network training, all kernels are random but static in time and the prediction variance has a universal form. However, in the rich, feature learning regime, the fluctuations of the kernels and predictions are dynamically coupled with a variance that can be computed self-consistently. In two layer networks, we show how feature learning can dynamically reduce the variance of the final tangent kernel and final network predictions. We also show how initialization variance can slow down online learning in wide but finite networks. In deeper networks, kernel variance can dramatically accumulate through subsequent layers at large feature learning strengths, but feature learning continues to improve the signal-to-noise ratio of the feature kernels. In discrete time, we demonstrate that large learning rate phenomena such as edge of stability effects can be well captured by infinite width dynamics and that initialization variance can decrease dynamically. For CNNs trained on CIFAR-10, we empirically find significant corrections to both the bias and variance of network dynamics due to finite width.

1 Introduction

Learning dynamics of deep neural networks are challenging to analyze and understand theoretically, but recent progress has been made by studying the idealization of infinite-width networks. Two types of infinite-width limits have been especially fruitful. First, the kernel or lazy infinite-width limit, which arises in the standard or neural tangent kernel (NTK) parameterization, gives prediction dynamics which correspond to a linear model [1-5]. This limit is theoretically tractable but fails to capture adaptation of internal features in the neural network, which are thought to be crucial to the success of deep learning in practice. Alternatively, the mean field or µ -parameterization allows feature learning at infinite width [6-9].

With a set of well-defined infinite-width limits, prior theoretical works have analyzed finite networks in the NTK parameterization perturbatively, revealing that finite width both enhances the amount of feature evolution (which is still small in this limit) but also introduces variance in the kernels and the predictions over random initializations [10-15]. Because of these competing effects, in some situations wider networks are better, and in others wider networks perform worse [16].

In this paper, we analyze finite-width network learning dynamics in the mean field parameterization. In this parameterization, wide networks are empirically observed to outperform narrow networks [7, 17, 18]. Our results and framework provide a methodology for reasoning about detrimental finitesize effects in such feature-learning neural networks. We show that observable averages involving kernels and predictions obey a well-defined power series in inverse width even in rich training regimes. We generally observe that the leading finite-size corrections to both the bias and variance components of the square loss are increased for narrower networks, and diminish performance. Further, we show that richer networks are closer to their corresponding infinite-width mean field limit. For simple tasks and architectures the leading O (1 / width ) corrections to the error can be descriptive, while for large sample size or more realistic tasks, higher order corrections appear to become relevant. Concretely, our contributions are listed below:

Starting from a dynamical mean field theory (DMFT) description of infinite-width nonlinear deep neural network training dynamics, we provide a complete recipe for computing fluctuation dynamics of DMFT order parameters over random network initializations during training. These include the variance of the training and test predictions and the O (1 / width ) variance of feature and gradient kernels throughout training.

We first solve these equations for the lazy limit, where no feature learning occurs, recovering a simple differential equation which describes how prediction variance evolves during learning.

We solve for variance in the rich feature learning regime in two-layer networks and deep linear networks. We show richer nonlinear dynamics improve the signal-to-noise ratio (SNR) of kernels and predictions, leading to closer agreement with infinite-width mean field behavior.

We analyze in a two-layer model why larger training set sizes in the overparameterized regime enhance finite-width effects and how richer training can reduce this effect.

We show that large learning rate effects such as edge-of-stability [19-21] dynamics can be well captured by infinite width theory, with finite size variance accurately predicted by our theory.

We test our predictions in Convolutional Neural Networks (CNNs) trained on CIFAR-10 [22]. We observe that wider networks and richly trained networks have lower logit variance as predicted. However, the timescale of training dynamics is significantly altered by finite width even after ensembling. We argue that this is due to a detrimental correction to the mean dynamical NTK.

1.1 Related Works

Infinite-width networks at initialization converge to a Gaussian process with a covariance kernel that is computed with a layerwise recursion [23-26, 13]. In the large but finite width limit, these kernels do not concentrate at each layer, but rather propagate finite-size corrections forward through the network [27-30, 14]. During gradient-based training with the NTK parameterization, a hierarchy of differential equations have been utilized to compute small feature learning corrections to the kernel through training [10-13]. However the higher order tensors required to compute the theory are initialization dependent, and the theory breaks down for sufficiently rich feature learning dynamics. Various works on Bayesian deep networks have also considered fluctuations and perturbations in the kernels at finite width during inference [31, 32]. Other relevant work in this domain are [33-39].

An alternative to standard/NTK parameterization is the mean field or µP -limit where features evolve even at infinite width [6-9, 40-42]. Recent studies on two-layer mean field networks trained online with Gaussian data have revealed that finite networks have larger sensitivity to SGD noise [43, 44]. Here, we examine how finite-width neural networks are sensitive to initialization noise. Prior work has studied how the weight space distribution and predictions converge to mean field dynamics with a dynamical error O (1 / √ width ) [40, 45], however in the deep case this requires a probability distribution over couplings between adjacent layers. Our analysis, by contrast, focuses on a function and kernel space picture which decouples interactions between layers at infinite width. A starting point for our present analysis of finite-width effects was a previous set of studies [9, 46] which identified the DMFT action corresponding to randomly initialized deep NNs which generates the distribution over kernel and network prediction dynamics. These prior works discuss the possibility of using a finite-size perturbation series but crucially failed to recognize the role of the network prediction fluctuations on the kernel fluctuations which are necessary to close the self-consistent equations in the rich regime. Using the mean field action to calculate a perturbation expansion around DMFT is a long celebrated technique to obtain finite size corrections in physics [47-50] and has been utilized for random untrained recurrent networks [51, 52], and more recently to calculate variance of

feature kernels Φ glyph[lscript] at initialization t = 0 in deep MLPs or RNNs [53]. We extend these prior studies to the dynamics of training and to probe how feature learning alters finite size corrections.

2 Problem Setup

We consider wide neural networks where the number of neurons (or channels for a CNN) N in each layer is large. For a multi-layer perceptron (MLP), the network is defined as a map from input x µ ∈ R D to hidden preactivations h glyph[lscript] µ ∈ R N in layers glyph[lscript] ∈ { 1 , ..., L } and finally output f µ

where γ is a scale factor that controls feature learning strength, with large γ leading to rich feature learning dynamics and the limit of small γ → 0 (or generally if γ scales as N -α for α > 0 as N →∞ , NTK parameterization corresponds to α = 1 2 ) gives lazy learning where no features are learned [4, 7, 9]. The parameters θ = { W 0 , W 1 , ..., w L } are optimized with gradient descent or gradient flow d dt θ = -Nγ 2 ∇ θ L where L = E x µ ∈D glyph[lscript] ( f ( x µ , θ ) , y µ ) is a loss computed over dataset D = { ( x 1 , y 1 ) , ( x 2 , y 2 ) , . . . ( x P , y P ) } . This parameterization and learning rate scaling ensures that d dt f µ ∼ O N,γ (1) and d dt h glyph[lscript] µ = O N,γ ( γ ) at initialization. This is equivalent to maximal update parameterization ( µ P)[8], which can be easily extended to other architectures including neural networks with trainable bias parameters as well as convolutional, recurrent, and attention layers [8, 9].

3 Review of Dynamical Mean Field Theory

The infinite-width training dynamics of feature learning neural networks was described by a DMFT in [9, 46]. We first review the DMFT's key concepts, before extending it to get insight into finite-widths. To arrive at the DMFT, one first notes that the training dynamics of such networks can be rewritten in terms of a collection of dynamical variables (or order parameters ) q = Vec { f µ ( t ) , Φ glyph[lscript] µν ( t, s ) , G glyph[lscript] µν ( t, s ) , ... } [9], which include feature and gradient kernels [9, 54]

where g glyph[lscript] µ ( t ) = γN ∂f µ ( t ) ∂ h glyph[lscript] µ ( t ) are the back-propagated gradient signals. Further, for widthN networks the distribution of these dynamical variables across weight initializations (from a Gaussian distribution θ ∼ N (0 , I ) ) is given by p ( q ) ∝ exp( NS ( q )) , where the action S ( q ) contains interactions between neuron activations and the kernels at each layer [9].

The DMFT introduced in [9] arises in the N → ∞ limit when p ( q ) is strongly peaked around the saddle point q ∞ where ∂S ∂ q | q ∞ = 0 . Analysis of the saddle point equations reveal that the training dynamics of the neural network can be alternatively described by a stochastic process. A key feature of this process is that it describes the training time evolution of the distribution of neuron pre-activations in each layer (informally the histogram of the elements of h glyph[lscript] µ ( t ) ) where each neuron's pre-activation behaves as an i.i.d. draw from this single-site stochastic process. We denote these random processes by h glyph[lscript] µ ( t ) . Kernels in (2) are now computed as averages over these infinite-width single site processes Φ glyph[lscript] µν ( t, s ) = 〈 φ ( h glyph[lscript] µ ( t )) φ ( h glyph[lscript] ν ( s )) 〉 , G glyph[lscript] µν ( t, s ) = 〈 g glyph[lscript] µ ( t ) g glyph[lscript] ν ( s ) 〉 , where averages arise from the N →∞ limit of the dot products in (2). DMFT also provides a set of self-consistent equations that describe the complete statistics of these random processes, which depend on the kernels, as well as other quantities. To make our notation and terminology clearer for a machine learning audience, we provide Table 1 for a definition of the physics terminology in machine learning language.

4 Dynamical Fluctuations Around Mean Field Theory

We are interested in going beyond the infinite-width limit to study more realistic finite-width networks. In this regime, the order parameters q fluctuate in a O (1 / √ N ) neighborhood of q ∞ [55, 51, 53, 46].

Statistics of these fluctuations can be calculated from a general cumulant expansion (see App. D) [55, 56, 51]. We will focus on the leading-order corrections to the infinite-width limit in this expansion.

Proposition 1 The finite-width N average of observable O ( q ) across initializations, which we denote by 〈 O ( q ) 〉 N , admits an expansion of the form whose leading terms are

where 〈〉 ∞ denotes an average over the Gaussian distribution q ∼ N ( q ∞ , -1 N ( ∇ 2 q S [ q ∞ ] ) -1 ) and the function V ( q ) ≡ S ( q ) -S ( q ∞ ) -1 2 ( q -q ∞ ) glyph[latticetop] ∇ 2 q S ( q ∞ )( q -q ∞ ) contains cubic and higher terms in the Taylor expansion of S around q ∞ . The terms shown include all the leading and sub-leading terms in the series in powers of 1 /N . The terms in ellipses are at least O ( N -1 ) suppressed compared to the terms provided.

The proof of this statement is given in App. D. The central object to characterize finite size effects is the unperturbed covariance (the propagator ): Σ = -[ ∇ 2 S ( q ∞ ) ] -1 . This object can be shown to capture leading order fluctuation statistics 〈 ( q -q ∞ ) ( q -q ∞ ) glyph[latticetop] 〉 N = 1 N Σ + O ( N -2 ) (App. D.1), which can be used to reason about, for example, expected square error over random initializations. Correction terms at finite width may give a possible explanation of the superior performance of wide networks at fixed γ [7, 17, 18]. To calculate such corrections, in App. E, we provide a complete description of Hessian ∇ 2 q S ( q ) and its inverse (the propagator) for a depthL network. This description constitutes one of our main results. The resulting expressions are lengthy and are left to App. E. Here, we discuss them at a high level. Conceptually there are two primary ingredients for obtaining the full propagator:

Hessian sub-blocks κ which describe the uncoupled variances of the kernels, such as

Similar terms also appear in other studies on finite width Bayesian inference [13, 31, 32] and in studies on kernel variance at initialization [27, 14, 29, 53].

Blocks which capture the sensitivity of field averages to pertubations of order parameters, such as

where ∆ µ ( t ) = -∂glyph[lscript] ( f µ ,y µ ) ∂f µ | f µ ( t ) are error signal for each data point.

Abstractly, we can consider the uncoupled variances κ as 'sources' of finite-width noise for each order parameter and the D blocks as summarizing a directed causal graph which captures how this noise propagates in the network (through layers and network predictions). In Figure 1, we illustrate this graph showing directed lines that represent causal influences of order parameters on fields and vice versa. For instance, if Φ glyph[lscript] were perturbed, D Φ glyph[lscript] +1 , Φ glyph[lscript] would quantify the resulting perturbation to Φ glyph[lscript] +1 through the fields h glyph[lscript] +1 .

In App. E, we calculate κ and D tensors, and show how to use them to calculate the propagator. As an example of our results:

Proposition 2 Partition q into primal q 1 = Vec { f µ ( t ) , Φ glyph[lscript] µν ( t, s ) ... } and conjugate variables q 2 = Vec { ˆ f µ ( t ) , ˆ Φ glyph[lscript] µν ( t, s ) ... } . Let κ = ∂ 2 ∂ q 2 ∂ q glyph[latticetop] 2 S [ q 1 , q 2 ] and D = ∂ 2 ∂ q 2 ∂ q glyph[latticetop] 1 S [ q 1 , q 2 ] , then the propagator for q 1 has the glyph[latticetop]

Φ

ℓ

𝐺

ℓ

Φ

ℓ"#

𝐺

ℓ"#

ℎ

ℓ

𝑔

ℓ

ℎ

ℓ"#

𝑔

ℓ"#

𝐾

Δ

OrderParameters

Fields

form Σ q 1 = D -1 κ [ D -1 ] (App E). The variables q 1 are related to network observables, while conjugates q 2 arise as Lagrange multipliers in the DMFT calculation. From the propagator Σ q 1 we can read off the variance of network observables such as N Var ( f µ ) ∼ Σ f µ .

The necessary order parameters for calculating the fluctuations are obtained by solving the DMFT using numerical methods introduced in [9]. We provide a pseudocode for this procedure in App. F. We proceed to solve the equations defining Σ in special cases which are illuminating and numerically feasible including lazy training, two layer networks and deep linear NNs.

5 Lazy Training Limit

To gain some initial intuition about why kernel fluctuations alter learning dynamics, we first analyze the static kernel limit γ → 0 where features are frozen. To prevent divergence of the network in this limit, we use a background subtracted function ˜ f ( x , θ ) = f ( x , θ ) -f ( x , θ 0 ) which is identically zero at initialization [4]. For mean square error, the N →∞ and γ → 0 limit is governed by ∂ ˜ f ( x ) ∂t = E x ′ ∼D ∆( x ′ ) K ( x , x ′ ) with ∆( x ) = y ( x ) -˜ f ( x ) (for MSE) and K is the static (finite width and random) NTK. The finiteN initial covariance of the NTK has been analyzed in prior works [27, 13, 14], which reveal a dependence on depth and nonlinearity. Since the NTK is static in the γ → 0 limit, it has constant initialization variance through training. Further, all sensitivity blocks of the Hessian involving the kernels and the prediction errors ∆ (such as the D Φ glyph[lscript] , ∆ ) vanish. We represent the covariance of the NTK as κ ( x 1 , x 2 , x 2 , x 3 ) = N Cov ( K ( x 1 , x 2 ) , K ( x 3 , x 4 )) . To identify the dynamics of the error ∆ covariance, we relate K , the finite width NTK, to K ∞ which is the deterministic infinite width NTK K ∞ . We consider the eigendecomposition of the infinite-width NTK K ∞ ( x , x ′ ) = ∑ k λ k ψ k ( x ) ψ k ( x ′ ) with respect to the training distribution D , and decompose κ in this basis.

where averages are computed over the training distribution D .

Proposition 3 For MSE loss, the prediction error covariance Σ ∆ ( t, s ) = N Cov 0 ( ∆ ( t ) , ∆ ( s )) satisfies a differential equation (App. H)

where ∆ ∞ k ( t ) ≡ exp( -λ k t ) 〈 ψ k ( x ) y ( x ) 〉 x are the errors at infinite width for eigenmode k .

0

50

100

150

200

t

1.0

0.5

0.0

0.5

1.0

(

t

),

(

t

)

(a) Average Prediction Errors

0

50

100

150

200

t

0.0

0.1

0.2

0.3

0.4

N

Var

k

(

t

)

k

=1

k

=2

k

=3

k

=4

k

=5

k

=6

k

=7

(b) Top Eigenmodes Variance

0

50

100

150

200

t

0.0

0.5

1.0

1.5

2.0

N

Var (

(

t

))

Ensemble (

N

=100)

DMFT

(c) Total Train Variance

glyph[negationslash]

An example verifying these dynamics is provided in Figure 2. In the case where the target is an eigenfunction y = ψ k ∗ , the covariance has the form Σ ∆ kglyph[lscript] ( t, s ) = κ kglyph[lscript]k ∗ k ∗ exp( -λ k ∗ ( t + s )) ( λ k -λ k ∗ )( λ glyph[lscript] -λ k ∗ ) . If the kernel is rank one with eigenvalue λ , then the dynamics have the simple form Σ ∆ ( t, s ) = κy 2 t s e -λ ( t + s ) . We note that similar terms appear in the prediction dynamics obtained by truncating the Neural Tangent Hierarchy [10, 11], however those dynamics concerned small feature learning

10 1

10 2

10 3

N

10

5

10

4

10

3

10

2

10

1

[

q

(

t

)

q

(

t

)]

2

K

K

f

(a) Deviation from DMFT vs

N

0

50

100

150

200

250

t

1.5

2.0

2.5

3.0

3.5

K

(

t

)

=0.50

=1.00

=2.00

=4.00

(b) Average NTK vs DMFT

0

50

100

150

200

250

t

0.0

0.1

0.2

0.3

0.4

0.5

0.6

f

(

t

)

(c) Average test prediction

0

50

100

150

200

250

t

0.00

0.02

0.04

0.06

0.08

0.10

N

Var(

)

=0.50

=1.00

=2.00

=4.00

(d) Train error variance

0

50

100

150

200

250

t

0.000

0.025

0.050

0.075

0.100

0.125

0.150

N

Var(

f

)

(e) Test Prediction Variance

0

50

100

150

200

250

t

1.00

1.25

1.50

1.75

2.00

2.25

N

Var

K

(

t

)

(f) Kernel Variance Dynamics

corrections rather than from initialization variance (App. H.1). Corrections to the mean 〈 ∆ 〉 are analyzed in App. H.2. We find that the variance and mean correction dynamics involves non-trivial coupling across eigendirections with a mixture of exponentials with timescales { λ -1 k } .

6 Rich Regime in Two-Layer Networks

In this section, we analyze how feature learning alters the variance through training. We show a denoising effect where the signal to noise ratios of the order parameters improve with feature learning.

6.1 Kernel and Error Coupled Fluctuations on Single Training Example

In the rich regime, the kernel evolves over time but inherits fluctuations from the training errors ∆ . To gain insight, we first study a simplified setting where the data distribution is a single training example x and single test point x glyph[star] in a two layer network. We will track ∆( t ) = y -f ( x , t ) and the test prediction f glyph[star] ( t ) = f ( x glyph[star] , t ) . To identify the dynamics of these predictions we need the NTK K ( t ) on the train point, as well as the train-test NTK K glyph[star] ( t ) . In this case, all order parameters can be viewed as scalar functions of a single time index (unlike the deep network case, see App. E).

Proposition 4 Computing the Hessian of the DMFT action and inverting (App. I), we obtain the following covariance for q 1 = Vec { ∆( t ) , f glyph[star] ( t ) , K ( t ) , K glyph[star] ( t ) } t ∈ R +

0

100

200

300

400

500

t

10

4

10

2

10 0

N

Var

f

N

=500

Var. Along

y

P

P

P

(a) Offline prediction Variance

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

Train Error

P

=10

N

=100

N

=100 | DMFT

N

=250

N

=1000

N

=

DMFT

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

P

=50

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

P

=100

(b) Offline leading correction breaks down for large

P

0

100

200

300

400

500

t

10

2

10

1

10 0

10 1

N

Var

f

N

=250,

P

=50

=1.0

=2.0

=3.0

(c) Large

γ

reduces variance

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

Train Error

P

=10

=1.0

=1.0 | DMFT

N

=

DMFT

=2.0

=3.0

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

P

=50

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

P

=100

(d) Large

γ

well captured by DMFT First Order Correction

0

100

200

300

400

500

t

10

4

10

2

10 0

N

i

Var

i

N

=500

Var. Along

D

D

D

(e) Online function variance

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

Generalization Error

D

=10

N

=100

N

=100 | DMFT

N

=250

N

=1000

N

=

DMFT

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

D

=50

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

D

=100

(f) Online leading correction breaks down for large

D

quantify sensitivity of the kernel to perturbations in the error signal ∆( s ) . Lastly κ and κ glyph[star] are the uncoupled variances of K ( t ) and K glyph[star]glyph[star] ( t ) and κ glyph[star] is the uncoupled covariance of K ( t ) , K glyph[star] ( t ) .

In Fig. 3, we plot the resulting theory (diagonal blocks of Σ q 1 from Equation 8) for two layer neural networks. As predicted by theory, all average squared deviations from the infinite width DMFT scale as O ( N -1 ) . Similarly, the average kernels 〈 K 〉 and test predictions 〈 f glyph[star] 〉 change by a larger amount for larger γ (equation (79)). The experimental variances also match the theory quite accurately. The variance of the train error ∆( t ) peaks earlier and at a lower value for richer training, but all variances go to zero at late time as the model approaches the interpolation condition ∆ = 0 . As γ → 0 the curve approaches N Var (∆( t )) ∼ κ y 2 t 2 e -2 t , where κ is the initial NTK variance (see Section 5). While the train prediction variance goes to zero, the test point prediction does not, with richer networks reaching a lower asymptotic variance. We suspect this dynamical effect could explain lower variance observed in feature learning networks compared to lazy networks [7, 18]. In Fig. A.1, we show that the reduction in variance is not due to a reduction in the uncoupled variance κ ( t, s ) , which increases in γ . Rather the reduction in variance is driven by the coupling of perturbations across time.

6.2 Offline Training with Multiple Samples or Online Training in High Dimension

In this section we go beyond the single sample equations of the prior section and explore training with multiple P examples. In this case, we have training errors { ∆ µ ( t ) } P µ =1 and multiple kernel entries K µν ( t ) (App. E). Each of the errors ∆ µ ( t ) receives a O ( N -1 / 2 ) fluctuation, the training error ∑ µ 〈 ∆ 2 µ 〉 has an additional variance on the order of O ( P N ) . In the case of two-layer linear

0

5

10

15

20

25

30

t

0.0

0.2

0.4

0.6

0.8

1.0

N

Var

(

t

)

= 0.5

= 1.0

= 2.0

= 3.0

(a) Prediction Variance

0

5

10

15

20

25

30

t

0.25

0.50

0.75

1.00

1.25

1.50

1.75

N

Var

K

(

t

) /

K

(

t

)

2

(b) NTK variance

0

10

20

30

5

10

15

N

Var

H

=1.0

=1

=2

=3

0

10

20

30

5

10

15

N

Var

G

=1

=2

=3

0

10

20

30

0

10

20

30

40

50

=2.0

0

10

20

30

0

10

20

30

40

50

0

10

20

30

0

50

100

150

=3.0

0

10

20

30

0

50

100

150

networks trained on whitened data ( 1 D x µ · x ν = δ µν ), the equations for the propagator simplify and one can separately solve for the variance of ∆ ( t ) ∈ R P along signal direction y ∈ R P and along each of the P -1 orthogonal directions (App. J). At infinite width, the task-orthogonal component ∆ ⊥ vanishes and only the signal dimension ∆ y ( t ) evolves in time with differential equation [9, 46]

However, at finite width, both the ∆ y ( t ) and the P -1 orthogonal variables ∆ ⊥ inherit initialization variance, which we represent as Σ ∆ y ( t, s ) and Σ ⊥ ( t, s ) . In Fig. 4 (a)-(b) we show this approximate solution 〈 | ∆ ( t ) | 2 〉 ∼ ∆ y ( t ) 2 + 2 N ∆ 1 y ( t )∆ y ( t ) + 1 N Σ ∆ y ( t, t ) + ( P -1) N Σ ⊥ ( t, t ) + O ( N -2 ) across varying γ and varying P (see Appendix J for Σ ∆ y and Σ ⊥ formulas). We see that variance of train point predictions f µ ( t ) increases with the total number of points despite the signal of the target vector ∑ µ y 2 µ being fixed. In this model, the bias correction 2 N ∆ 1 y ( t )∆ y ( t ) is always O (1 /N ) but the variance correction is O ( P/N ) . The fluctuations along the P -1 orthogonal directions begin to dominate the variance at large P . Fig. 4 (b) shows that as P increases, the leading order approximation breaks down as higher order terms become relevant. Analysis for online training reveals identical fluctuation statistics, but with variance that scales as ∼ D/N (Appendix K) as we verify in Figure 4 (e)-(f).

7 Deep Networks

In networks deeper than two layers, the DMFT propagator has complicated dependence on nondiagonal (in time) entries of the feature kernels (see App. E). This leads to Hessian blocks with four time and four sample indices such as D Φ glyph[lscript] µναβ ( t, s, t ′ , s ′ ) = ∂ ∂ Φ glyph[lscript] -1 αβ ( t ′ ,s ′ ) 〈 φ ( h glyph[lscript] µ ( t )) φ ( h glyph[lscript] ν ( s )) 〉 , rendering any numerical calculation challenging. However, in deep linear networks trained on whitened data, we can exploit the symmetry in sample space and the Gaussianity of preactivation features to exactly compute derivatives without Monte Carlo sampling as we discuss in App. L. An example set of results for a depth 4 network is provided in Fig. 5. The variance for the feature kernels H glyph[lscript] accumulate finite size variance by layer along the forward pass and the gradient kernels G glyph[lscript] accumulate variance on the backward pass. The SNR of the kernels 〈 H 〉 2 N Var ( H ) improves with feature learning, suggesting that richer networks will be better modeled by their mean field limits. Examples of the off-diagonal correlations obtained from the propagator are provided in App. Fig. A.3.

8 Variance can be Small Near Edge of Stability

In this section, we move beyond the gradient flow formalism and ask what large step sizes do to finite size effects. Recent studies have identified that networks trained at large learning rates can be qualitatively different than networks in the gradient flow regime, including the catapult [57] and edge of stability (EOS) phenomena [19-21]. In these settings, the kernel undergoes an initial scale

0

5

10

15

20

25

30

t

0.25

0.50

0.75

1.00

1.25

1.50

2

K

(

t

)

(a) Mean Kernel Dynamics

0

5

10

15

20

25

30

t

10

5

10

4

10

3

10

2

10

1

10 0

N

Var

(

t

)

(b) Error Variance

0

5

10

15

20

25

30

t

10

4

10

3

10

2

10

1

10 0

10 1

10 2

N

Var

K

(

t

)

growth before exhibiting either a recovery or a clipping effect. In this section, we explore whether these dynamics are highly sensitive to initialization variance or if finite networks are well captured by mean field theory. Following [57], we consider two layer networks trained on a single example | x | 2 = D and y = 1 . We use learning rate η and feature learning strength γ . The infinite width mean field equations for the prediction f t and the kernel K t are (App. M)

For small η , the equations are well approximated by the gradient flow limit and for small γ corresponds to a discrete time linear model. For large ηγ > 1 , the kernel K progressively sharpens (increases in scale) until it reaches 2 /η and then oscillates around this value. It may be expected that near the EOS, the large oscillations in the kernels and predictions could lead to amplified finite size effects, however, we show in Fig. 6 that the leading order propagator elements decrease even after reaching the EOS threshold, indicating reduced disagreement between finite and infinite width dynamics.

9 Finite Width Alters Bias, Training Rate, and Variance in Realistic Tasks

To analyze the effect of finite width on neural network dynamics during realistic learning tasks, we studied a vanilla depth6 ReLU CNN trained on CIFAR-10 (experimental details in App. B, G.2) In Fig. 7, we train an ensemble of E = 8 independently initialized CNNs of each width N . Wider networks not only have better performance for a single model (solid), but also have lower bias (dashed), measured with ensemble averaging of the logits. Because of faster convergence of wide networks, we observe wider networks have higher variance, but if we plot variance at fixed ensembled training accuracy, wider networks have consistently lower variance (Fig. 7(d)).

We next seek an explanation for why wider networks after ensembling trains at a faster rate . Theoretically, this can be rationalized by a finite-width alteration to the ensemble averaged NTK, which governs the convergence timescale of the ensembled predictions (App. G.1). Our analysis in App. G.1 suggests that the rate of convergence receives a finite size correction with leading correction O ( N -1 ) G.2. To test this hypothesis, we fit the ensemble training loss curve to exponential function L ≈ C exp( -R N t ) where C is a constant. We plot the fit R N as a function of N -1 result in Fig. 7(e). For large N , we see the leading behavior is linear in N -1 , but begins to deviate at small N as a quadratic function of N -1 , suggesting that second order effects become relevant around N glyph[lessorsimilar] 100 .

In App. Fig. A.4, we train a smaller subset of CIFAR-10 where we find that R N is well approximated by a O ( N -1 ) correction, consistent with the idea that higher sample size drives the dynamics out of the leading order picture. We also analyze the effect of γ on variance in this task. In App. Fig. A.5, we train N = 64 models with varying γ . Increased γ reduces variance of the logits and alters the representation (measured with kernel-task alignment), the training and test accuracy are roughly insensitive to the richness γ in the range we considered.

10 Discussion

We studied the leading order fluctuations of kernels and predictions in mean field neural networks. Feature learning dynamics can reduce undesirable finite size variance, making finite networks order

0

20

40

60

80

100

Epochs

0.01

0.02

0.03

0.04

Train Loss

N

=32

N

=45

N

=64

N

=90

N

=128

N

=256

(a) Train Loss Across Widths

0

20

40

60

80

100

Epochs

0.2

0.4

0.6

Test Acc.

N

=32

N

=45

N

=64

N

=90

N

=128

N

=256

(b) Test Acc. Across Widths

0

20

40

60

80

100

Epochs

0.02

0.04

0.06

0.08

0.10

A

(

L

,

yy

)

N

=32

N

=45

N

=64

N

=90

N

=128

N

=256

(c) Alignment Dynamics

0.2

0.4

0.6

0.8

1.0

Ens. Train Acc

1.00

1.05

1.10

1.15

1.20

Test MSE Ratio

N

=32

N

=45

N

=64

N

=90

N

=128

N

=256

(d) Test MSE Ensembling Ratio

0.00

0.01

0.02

0.03

N

1

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

0.0175

RN

0.017

0.394/

N

0.017

0.430/

N

6.929/

N

2

(e) Training Rates vs

N

0.0

0.5

1.0

1.5

Rescaled Time

0.01

0.02

0.03

0.04

Train Loss

N

=32

N

=45

N

=64

N

=90

N

=128

N

=256

(f) Train Acc. Rescaled Time

parameters closer to the infinite width limit. In several toy models, we revealed some interesting connections between the influence of feature learning, depth, sample size, and large learning rate and the variance of various DMFT order parameters. Lastly, in realistic tasks, we illustrated that bias corrections can be significant as rates of learning can be modified by width. Though our full set of equations for the leading finite size fluctuations are quite general in terms of network architecture and data structure, they are only derived at the level of rigor of physics rather than a formally rigorous proof which would need several additional assumptions to make the perturbation expansion properly defined. Further, the leading terms in our perturbation series involving only Σ does not capture the complete finite size distribution defined in Eq. (3), especially as the sample size becomes comparable to the width. It would be interesting to see if proportional limits of the rich training regime where samples and width scale linearly can be examined dynamically [58]. Future work could explore in greater detail the higher order contributions from averages involving powers of V ( q ) by examining cubic and higher derivatives of S in Eq. (3). It could also be worth examining in future work how finite size impacts other biologically plausible learning rules, where the effective NTK can have asymmetric (over sample index) fluctuations [46]. Also of interest would be computing the finite width effects in other types of architectures, including residual networks with various branch scalings [59, 60]. Further, even though we expect our perturbative expressions to give a precise asymptotic description of finite networks in mean field/ µ P, the resulting expressions are not realistically computable in deep networks trained on large dataset size P for long times T since the number of Hessian entries scales as O ( T 4 P 4 ) and a matrix of this size must be stored in memory and inverted in the general case. Future work could explore solveable special cases such as high dimensional limits.

Code Availability

Code to reproduce the experiments in this paper is provided at https://github.com/ Pehlevan-Group/dmft_fluctuations . Details about numerical methods and computational implementation can be found in Appendices F and N.

References

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31, pages 8571-8580. Curran Associates, Inc., 2018.

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems , 32, 2019.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems , 32, 2019.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems , 32, 2019.

Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. In International Conference on Machine Learning , pages 1176211772. PMLR, 2021.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , pages 2388-2464. PMLR, 2019.

Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment , 2020(11):113301, 2020.

Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning , pages 11727-11737. PMLR, 2021.

Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. In International conference on machine learning , pages 4542-4551. PMLR, 2020.

Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations , 2020.

Anders Andreassen and Ethan Dyer. Asymptotics of wide convolutional neural networks. arXiv preprint arXiv:2008.08675 , 2020.

Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory . Cambridge University Press Cambridge, MA, USA, 2022.

Boris Hanin. Random fully connected neural networks as perturbatively solvable hierarchies. arXiv preprint arXiv:2204.01058 , 2022.

Sho Yaida. Meta-principled family of hyperparameter scaling strategies. arXiv preprint arXiv:2210.04909 , 2022.

Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems , 33:15156-15172, 2020.

Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems , 34, 2021.

Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. In The Eleventh International Conference on Learning Representations , 2023.

Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations , 2021.

Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations , 2023.

Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. arXiv preprint arXiv:2210.04860 , 2022.

Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research).

Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks , pages 29-53. Springer, 1996.

Radford M Neal. Bayesian learning for neural networks , volume 118. Springer Science & Business Media, 2012.

Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In International Conference on Learning Representations , 2018.

Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations , 2018.

Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations , 2020.

Boris Hanin and Mihai Nica. Products of many large random matrices and gradients in deep neural networks. Communications in Mathematical Physics , 376(1):287-322, 2020.

Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning , pages 165-192. PMLR, 2020.

Jacob Zavatone-Veth and Cengiz Pehlevan. Exact marginal prior distributions of finite bayesian neural networks. Advances in Neural Information Processing Systems , 34:3364-3375, 2021.

Jacob Zavatone-Veth, Abdulkadir Canatar, Ben Ruben, and Cengiz Pehlevan. Asymptotics of representation learning in finite bayesian neural networks. Advances in Neural Information Processing Systems , 34, 2021.

Gadi Naveh, Oded Ben David, Haim Sompolinsky, and Zohar Ringel. Predicting the outputs of finite deep neural networks trained with noisy gradients. Physical Review E , 104(6):064301, 2021.

Adam X. Yang, Maxime Robeyns, Edward Milsom, Ben Anson, Nandi Schoots, and Laurence Aitchison. A theory of representation learning gives a deep generalisation of kernel methods. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 39380-39415. PMLR, 23-29 Jul 2023.

Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some cnns. Nature Communications , 14(1):908, 2023.

Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X , 11(3):031059, 2021.

Jacob A Zavatone-Veth and Cengiz Pehlevan. Depth induces scale-averaging in overparameterized linear bayesian neural networks. 55th Asilomar Conference on Signals, Systems, and Computers , 2021.

Jacob A Zavatone-Veth, William L Tong, and Cengiz Pehlevan. Contrasting random and learned features in deep bayesian linear regression. Physical Review E , 105(6):064118, 2022.

Gadi Naveh and Zohar Ringel. A self consistent theory of gaussian processes captures feature learning effects in finite cnns. Advances in Neural Information Processing Systems , 34, 2021.

Jacob A Zavatone-Veth, Abdulkadir Canatar, Benjamin S Ruben, and Cengiz Pehlevan. Asymptotics of representation learning in finite bayesian neural networks*. Journal of Statistical Mechanics: Theory and Experiment , 2022(11):114008, nov 2022.

Grant Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems , 31, 2018.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

Dyego Araújo, Roberto I Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.00193 , 2019.

Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborova. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From highdimensional amp; mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirty Sixth Conference on Learning Theory , volume 195 of Proceedings of Machine Learning Research , pages 1199-1227. PMLR, 12-15 Jul 2023.

Huy Tuan Pham and Phan-Minh Nguyen. Limiting fluctuation and trajectorial stability of multilayer neural networks with mean field training. Advances in Neural Information Processing Systems , 34:4843-4855, 2021.

Blake Bordelon and Cengiz Pehlevan. The influence of learning rule on representation dynamics in wide neural networks. In The Eleventh International Conference on Learning Representations , 2023.

Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A , 8(1):423, 1973.

Moshe Moshe and Jean Zinn-Justin. Quantum field theory in the large n limit: A review. Physics Reports , 385(3-6):69-228, 2003.

Jean Zinn-Justin. Quantum field theory and critical phenomena , volume 171. Oxford university press, 2021.

Carson C Chow and Michael A Buice. Path integral methods for stochastic differential equations. The Journal of Mathematical Neuroscience (JMN) , 5:1-35, 2015.

Moritz Helias and David Dahmen. Statistical Field Theory for Neural Networks . Springer International Publishing, 2020.

A Crisanti and H Sompolinsky. Path integral approach to random neural networks. Physical Review E , 98(6):062120, 2018.

Kai Segadlo, Bastian Epping, Alexander van Meegen, David Dahmen, Michael Krämer, and Moritz Helias. Unified field theoretical approach to deep and recurrent neuronal networks. Journal of Statistical Mechanics: Theory and Experiment , 2022(10):103401, 2022.

Yizhang Lou, Chris E Mingard, and Soufiane Hayou. Feature learning and signal propagation in deep neural networks. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 1424814282. PMLR, 17-23 Jul 2022.

Carl M Bender and Steven Orszag. Advanced mathematical methods for scientists and engineers I: Asymptotic methods and perturbation theory , volume 1. Springer Science & Business Media, 1999.

Mehran Kardar. Statistical physics of fields . Cambridge University Press, 2007.

Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218 , 2020.

Ben Adlam and Jeffrey Pennington. The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization. In International Conference on Machine Learning , pages 74-84. PMLR, 2020.

Soufiane Hayou. On the infinite-depth limit of finite-width neural networks. Transactions on Machine Learning Research , 2023.

Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309.16620 , 2023.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023.

Michael Baake and Ulrike Schlaegel. The peano-baker series. Proceedings of the Steklov Institute of Mathematics , 275(1):155-159, 2011.

Roger W Brockett. Finite dimensional linear systems . SIAM, 2015.

Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations , 2022.

B. Bordelon, A. Canatar, and C. Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. International Conference of Machine Learning , 2020.

Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint arXiv:1902.02880 , 2019.

Appendix

A Additional Figures

0

50

100

150

200

250

t

2.0

2.5

3.0

3.5

(

t

)

=0.5

=1.0

=2.0

=4.0

(a) Fourth Moment

κ

(

t, t

)

0

50

100

150

200

250

t

0.00

0.02

0.04

0.06

D

(

t

, 100)

(b) Kernel sensitivity

D

(

t,

100)

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

N

i

Var

i

N

=500,

D

=100

Expt

DMFT Limit

Leading Var.

Leading Var. + Bias

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

N

=500,

D

=500

0

100

200

300

400

500

t

10

5

10

4

10

3

10

2

10

1

10 0

N

=500,

D

=800

0

10

20

0

10

20

Expt

H

1 (

t

,

s

)

=0.5

1.95

1.96

1.97

1.98

1.99

2.00

0

10

20

0

10

20

DMFT

H

1 (

t

,

s

)

1.99

2.00

2.01

2.02

2.03

0

10

20

0

10

20

=1.0

1.90

1.95

2.00

2.05

2.10

0

10

20

0

10

20

1.95

2.00

2.05

2.10

0

10

20

0

10

20

=2.0

1.8

2.0

2.2

2.4

0

10

20

0

10

20

1.8

2.0

2.2

2.4

0

10

20

0

10

20

=3.0

1.75

2.00

2.25

2.50

2.75

0

10

20

0

10

20

1.75

2.00

2.25

2.50

2.75

3.00

0

10

20

0

10

20

Expt

H

2 (

t

,

s

)

=0.5

4.0

4.1

4.2

4.3

4.4

4.5

0

10

20

0

10

20

DMFT

H

2 (

t

,

s

)

4.0

4.1

4.2

4.3

4.4

4.5

0

10

20

0

10

20

=1.0

4.0

4.5

5.0

5.5

6.0

0

10

20

0

10

20

4.0

4.5

5.0

5.5

0

10

20

0

10

20

=2.0

4

6

8

10

0

10

20

0

10

20

4

6

8

10

0

10

20

0

10

20

=3.0

5.0

7.5

10.0

12.5

15.0

17.5

20.0

0

10

20

0

10

20

5

10

15

20

0

10

20

0

10

20

Expt

H

3 (

t

,

s

)

=0.5

6.0

6.5

7.0

7.5

8.0

8.5

0

10

20

0

10

20

DMFT

H

3 (

t

,

s

)

6.0

6.5

7.0

7.5

8.0

0

10

20

0

10

20

=1.0

6

8

10

12

14

16

0

10

20

0

10

20

6

8

10

12

14

0

10

20

0

10

20

=2.0

10

20

30

40

50

0

10

20

0

10

20

10

20

30

40

50

0

10

20

0

10

20

=3.0

20

40

60

80

100

120

0

10

20

0

10

20

20

40

60

80

100

120

140

0

100

200

300

Steps

2×10 2

3×10 2

4×10 2

Ens. Train Loss

N = 32

exp(

R

32

t

)

N = 64

N = 128

N = 256

N = 512

(a) Ens. Averaged Train Loss

0

100

200

300

Steps

10 0

10 1

10 2

10 3

10 4

K

|

|

2

(b) Raleigh Quotient

∆

glyph[latticetop]

K

∆

|

∆

|

2

0.000 0.005 0.010 0.015 0.020 0.025 0.030

N

1

0.0020

0.0025

0.0030

0.0035

RN

0.0035

0.0581/

N

(c) Rates as Function of

N

0

50

100

150

200

250

Epochs

0.02

0.03

0.04

0.05

Train Loss

=0.035

=0.050

=0.100

=0.200

=0.300

(a) Training MSE

0

50

100

150

200

250

Epochs

0.02

0.04

0.06

0.08

0.10

A

(

L

,

yy

)

(b) Kernel-Task Alignment

0

50

100

150

200

250

Epochs

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Test Acc

(c) Test Acc.

0

50

100

150

200

250

Epochs

1.00

1.05

1.10

1.15

1.20

Train MSE Ratio

(d) Single vs Ensembled Ratio

0

50

100

150

200

250

Epochs

1.00

1.02

1.04

1.06

1.08

1.10

Test MSE Ratio

(e) Single vs Ensembled Ratio

B CIFAR-10 Experimental Details

We trained the following depth 6 CNN architecture in the mean field parameterization using FLAX [61] on a single GPU. The bias parameters were zero in each hidden Conv layer, but were used for the readout weights. The networks were trained with MSE loss on centered 10 dimensional targets y µ ∈ R 10 for µ ∈ [ P ] . Each convolution was followed by an average pooling operation. To obtain mean field behavior, NTK parameterization with a modified final layer is used [7, 9].

All models were trained with standard SGD with a batch size of 256 . Each element in the ensemble of E networks is trained on identical batches presented in identical order. For the Figure 7 experiments, the raw learning rate is scaled as η = 10 N √ γ with γ = 0 . 2 (note that mean field theory requires scaling the raw learning rate linearly with N since the raw NTK is O ( N -1 ) [9]). For Figure A.5, the learning rate is η = 5 N √ γ . We find that choosing η ∝ √ γ gives approximately conserved training

times across γ (though distinct representation dynamics). The Figure A.4 shows the dynamics of fitting P = 64 training points with full batch gradient descent and γ = 0 . 1 .

C Review of DMFT: Deriving the Action

In this section we derive the DMFT action which contains all of the necessary statistical information about randomly initialized finite width N networks. From the action S the DMFT saddle point and the propagator can be computed. This derivation follows closely the original derivation by Bordelon &Pehlevan [9]. We start by writing the gradient flow dynamics on weight matrices

where ∆ µ ( t ) = -∂ L ∂f µ ( t ) are the error signals and g glyph[lscript] µ ( t ) = Nγ ∂f µ ( t ) ∂ h glyph[lscript] µ ( t ) are the back-propagation signals. The prediction dynamics satisfy

where K µν ( t ) is the instantaneous neural tangent kernel (NTK). At finite width N all of the above quantities depend on the precise initialization of the network. We transform the weight dynamics into an integral equation and use the recurrence for h glyph[lscript] to obtain the following

where we introduced the feature and gradient kernels

Written this way, we see that the source of the disorder which depends on the initial random weights W glyph[lscript] (0) comes through the fields

If we can characterize the distribution of the fields χ glyph[lscript] µ ( t ) and ξ glyph[lscript] µ ( t ) , then we can consequently characterize the distribution of h glyph[lscript] µ ( t ) , g glyph[lscript] µ ( t ) . We therefore choose to study the moment generating functional

Moments of these fields can be computed through differentiation with respect to the sources j , v near zero-source ( j = v = 0 )

To average over the initial weights, we introduce a Fourier representation of the Dirac-Delta function 1 = ∫ dzδ ( z ) = ∫ dzd ˆ z 2 π exp( i ˆ zz ) . We perform this transformation for each of the fields to enforce

their definition

We insert these Dirac delta functions so that we can directly average over the weights

where we introduced the kernels Φ glyph[lscript] , G glyph[lscript] . We next introduce the order parameter

To enforce the definitions of the new order parameters { Φ , G, A } we again introduce Dirac-delta functions

Analogous constraints for G and A are enforced with conjugate variables ˆ G,B . After introducing these variables, we find that the moment generating functional has the form

where S is the O (1) DMFT action which defines the statistical distribution over the dynamics. The action takes the form

where Z glyph[lscript] is the single site stochastic process for layer glyph[lscript] which defines the marginal distribution of χ, ξ , with the following form

where in the above, the { h, g } fields should be regarded as functionals of { χ, ξ } . At zero source j glyph[lscript] , v glyph[lscript] → 0 this function S can be regarded as the log density for the complete collection of order parameters q = { ˆ Φ , Φ , ˆ G,G,A,B } which collectively control the dynamics. Concretely, we have that p ( q ) ∝ exp( NS ( q )) . In the next section we explore an approximation scheme for averages over this distribution at large N .

D Cumulant Expansion of Observables

We are interested in a principled power series expansion (in 1 /N ) of any observable average 〈 O ( q ) 〉 that depends on DMFT order parameters q . At any width N the observable average takes the form

As discussed in the main text, the N → ∞ limit gives 〈 O ( q ) 〉 N ∼ O ( q ∞ ) where ∂S ∂ q | q ∞ = 0 by a steepest descent argument [55]. We assume that S 's Hessian is negative semidefinite so that Σ ≡ -[ ∇ 2 S ( q ) | q ∞ ] -1 glyph[followsequal] 0 and Taylor expand S ( q ) around the saddle point q ∞ giving S ( q ) = S ( q ∞ ) + 1 2 ( q -q ∞ ) glyph[latticetop] ∇ 2 S ( q )( q -q ∞ ) + V ( q -q ∞ ) . We note that the remainder function V contains only cubic and higher powers of q -q ∞ ≡ δ / √ N . The variable δ will be order O (1) . This will allow us to verify that additional terms are suppressed in powers of 1 /N . Expanding both the numerator and denominator's integrands in powers of V , we find

where 〈〉 ∞ represents an average over the Gaussian fluctuation N ( q ∞ , -1 N [ ∇ 2 q S ( q ∞ ) ] -1 ) . We see that the series in the denominator contains terms of the form N k k ! 〈 V k 〉 ∞ while the numerator depends on terms of the form N k k ! 〈 V k O 〉 ∞ / 〈 O 〉 ∞ . In either of these power series, the k -th term can contribute at most

since V contributes only cubic and higher terms. Thus each term in the numerator and denominator's series contains increasing powers of 1 /N . Concretely, each of the two series have terms of order { N 0 , N -1 , N -1 , N -2 , N -2 , ... } . Thus any quantity of the form 〈 O 〉 〈 O 〉 ∞ admits a ratio of power series in powers of 1 /N . One could truncate each of the series in the numerator and denominator to a desired order in N . Alternatively, the denominator could be expanded giving a single series (the cumulant expansion [56]). The first few terms in the cumulant expansion have the form

In this work, we mainly are interested in the leading order correction to 〈 O 〉 which can always be obtained with the truncation after the terms linear in V for any observable O .

D.1 Square Deviation from DMFT

We will now analyze the fluctuation statistics of our order parameters around the saddle point 〈 ( q -q ∞ )( q -q ∞ ) glyph[latticetop] 〉 N which has the form

as stated in the main text and verified empirically in Figure 3 (a). The reason that the terms in the numerator involving V can be no larger than O ( N -2 ) comes from vanishing of odd moments for q -q ∞ in the unperturbed distribution. Thus the leading expression for 〈 ( q -q ∞ )( q -q ∞ ) glyph[latticetop] 〉 only depends on Σ and not on V .

D.2 Mean Deviation from DMFT

Although the square displacement from DMFT only depended on Σ and not on V , we note that the average order parameter displacement 〈 q -q ∞ 〉 does receive a O (1 /N ) correction that depends on the perturbed potential V

where in the last line we used Stein's lemma (Gaussian integration by parts) for the Gaussian distribution over q . Note that 〈 ∂V ∂ q 〉 ∞ ∼ O ( 1 N ) since the derivative of the cubic term in V gives a quadratic function of q -q ∞ , whose average must be O ( N -1 ) . In this work, we focus primarily on the structure of the propagator, but outline a general recipe for getting the leading mean correction in Appendix G and H.2.

D.3 Covariance of Order Parameters

Lastly, we combine the previous two observations to reason about the scaling of the order parameter covariance over initializations. We note that the leading covariance of the order parameters over random initializations is also given by the propagator: Cov ( q ) ∼ 1 N Σ + O ( N -2 ) , since

due to the arguments above which showed that 〈 ( q -q ∞ )( q -q ∞ ) glyph[latticetop] 〉 ∼ 1 N Σ + O ( N -2 ) and that q ∞ -〈 q 〉 N ∼ O ( N -1 ) . Therefore, in the leading order picture, it is safe to associate Σ with the covariance of order parameters over random initializations of the network weights.

E Propagator Structure for the full DMFT Action

In this section, we examine the propagator structure for the full DMFT action. This action is modified from other prior works [9, 46] to include the evolution of the network prediction errors ∆( t ) . Those prior works noted that ∆ and the NTK K are deterministic functions of deterministic order parameters { Φ glyph[lscript] , G glyph[lscript] } in the N →∞ limit so those authors did not explicitly include ∆ or K in the action. At finite width N , including ∆ , K in the action is crucial as the fluctuation in prediction errors ∆ has significant consequences for dynamical fluctuations of kernels through the preactivation and

pre-gradient fields. In this section, we will mainly focus on gradient flow, but we describe large step size in Appendix M.

where the single site moment generating functionals Z glyph[lscript] have the form

with g glyph[lscript] µ ( t ) = ˙ φ ( h glyph[lscript] µ ( t )) z glyph[lscript] µ ( t ) . The saddle point equations give the infinite width evolution of our order parameters.

These equations exactly recover the mean field description obtained [9]. Note that 〈〉 for field averages is an average defined by Z glyph[lscript] and is distinct from the types averages 〈〉 , 〈〉 ∞ we have been considering over the order parameters q . The complementary set of equations for the primal variables, such as ∂S ∂ Φ glyph[lscript] µν ( t,s ) = 0 , give that ˆ K = ˆ ∆ = ˆ Φ = ˆ G = 0 at the saddle point. We now set out to compute the Hessan ∇ 2 q S . To simplify the set of expressions, we will only explicitly write out the nonvanishing

blocks. We will start with second derivatives involving only pairs of dual variables { ˆ Φ , ˆ G,A,B }

Next, we consider the second derivatives involving only primal variables { Φ glyph[lscript] , G glyph[lscript] , K, ∆ } which all vanish

Now we consider all derivatives which involve one of the dual variables { ˆ Φ glyph[lscript] , ˆ G glyph[lscript] , A glyph[lscript] , B glyph[lscript] } and the primal variable ∆

Now, we consider the second derivatives involving one derivative on a dual variable { ˆ Φ glyph[lscript] , ˆ G glyph[lscript] , A, B } and one of the primal variables { Φ glyph[lscript] , G glyph[lscript] } .

We note that terms such as ∂ ∂ Φ glyph[lscript] -1 αβ ( t ′ ,s ′ ) 〈 φ ( h glyph[lscript] µ ( t )) φ ( h glyph[lscript] ν ( s )) 〉 can be further decomposed since the average over the { u glyph[lscript] µ ( t ) } ∼ GP (0 , Φ glyph[lscript] -1 ) and h glyph[lscript] 's explicit dynamics both depend on Φ glyph[lscript] -1

where the first term comes from differentiating the Gaussian probability density for u glyph[lscript] (e.g. Price's theorem) and the second term is an explicit derivative of the preactivation fields with u glyph[lscript] treated as constant. Next we consider the nonvanishing terms which involve { ˆ ∆ , ˆ K, ∆ , K } which give

This enumerates all possible non-vanishing terms in the Hessian. We can now construct a block matrix of these Hessians by partitioning our order parameters q = [ q 1 , q 2 ] glyph[latticetop] where

This choice will become apparent shortly.

To calculate the full propagator Σ = -[ ∇ 2 q S ] -1 , we will assume invertibility of the upper block Σ 0 = -[ ∇ 2 q 1 S ] -1 and use this in the Schur complement

We now need to solve for Σ 0 = -[ ∇ 2 q 1 S ] -1 . To perform this inverse, we again partition q 1 into two sets of order parameters q 1 = [ q 1 1 , q 2 1 ] where q 1 1 = Vec { Φ glyph[lscript] µν ( t, s ) , G glyph[lscript] µν ( t, s ) , K µν ( t ) , ∆ µ ( t ) } and q 2 1 = Vec { ˆ Φ glyph[lscript] µν ( t, s ) , ˆ G glyph[lscript] µν ( t, s ) , ˆ K µν ( t ) , ˆ ∆ µ ( t ) }

We seek a physically sensible inverse where the variance of q 2 1 is vanishing [51, 53]. This leads to the following sub-propagator Σ 0

Thus given κ , U , we can solve for Σ 0 and ultimately for the full propagator Σ . The relevant entries in κ and U are given by those second derivatives calculated above. We note that each of the field derivatives needed for U can be computed implicitly from the field dynamics. For example, for the ∆ µ ( t ) derivatives we have

These can then be used in the averages such as 〈 ∂ ∂ ∆ ν ′ ( t ′ ) φ ( h glyph[lscript] µ ( t )) φ ( h glyph[lscript] ν ( s )) 〉 . Similarly, we can compute terms such as ∂h glyph[lscript] µ ( t ) ∂ Φ glyph[lscript] αβ ( t ′ ,s ′ ) through the following closed equations

These terms can then be used to compute quantities like D Φ glyph[lscript] .

F Solving for the Propagator

In this section we sketch out the required steps to obtain the propagator Σ .

Step 1: Solve the infinite width DMFT equations for q ∞ which include the prediction error dynamics ∆ µ ( t ) , the feature kernels Φ glyph[lscript] µν ( t, s ) , gradient kernels G glyph[lscript] µν ( t, s ) . This step corresponds to algorithm in Bordelon & Pehlevan '22 and defines the dynamics one would expect at infinite width [9]. See below for more detail.

Step 2: Compute the entries of the Hessian of S evaluated at the q ∞ computed in the first step. Some of these entries look like fourth cumulants of features like κ = 〈 φ ( h ) 4 〉 -〈 φ ( h ) 2 〉 2 and some of them measure sensitivity of one order parameter to a perturbation in another order parameter D Φ glyph[lscript] = ∂ ∂ Φ glyph[lscript] -1 〈 φ ( h glyph[lscript] ) 2 〉 . The averages 〈〉 used to calculate κ and D Φ glyph[lscript] should be performed over the infinite width stochastic processes for preactivations h glyph[lscript] which are defined in equation (19).

Step 3: After populating the entries of the block matrix for the Hesssian ∇ 2 S , we then calculate the propagator Σ with a matrix inversion. Since we discretized time, this is a finite dimensional matrix.

The step 1 above demands a solution to the infinite width DMFT equations (solving for the saddle point q ∞ ). We will now give a detailed set of instructions about how the infinite width limit for q ∞ is solved (step 1 above). This corresponds to the algorithm of Bordelon & Pehlevan 2022 to solve the saddle point equations ∂ ∂ q S ( q ) | q ∞ = 0 [9].

Step 1: Start with a guess for the kernels Φ glyph[lscript] µν ( t, s ) , G glyph[lscript] µν ( t, s ) and for the predictions through time f µ ( t ) . We usually use the lazy limit (e.g. Φ glyph[lscript] µν ( t, s ) = Φ glyph[lscript] µν (0 , 0) ...) as an initial guess.

Step 2: Sample Gaussian sources u glyph[lscript] µ ( t ) and r glyph[lscript] µ ( t ) based on the current covariances Φ glyph[lscript] and G glyph[lscript] .

Step 3: For each sample, solve integral equations for h ( t ) and z ( t ) .

These will be samples from the single site distribution for h, z

Step 4: Average over the Monte Carlo samples to produce a new estimate of the kernels: Φ glyph[lscript] ( t, s ) = 〈 φ ( h glyph[lscript] ( t )) φ ( h glyph[lscript] ( s )) 〉 . A similar procedure is performed for G glyph[lscript] and the response functions A glyph[lscript] , B glyph[lscript] .

Step 5: Compute the NTK estimate K ( t ) = ∑ glyph[lscript] G glyph[lscript] +1 ( t, t )Φ glyph[lscript] ( t, t ) and then integrate prediction dynamics from the dynamics of the NTK d dt f µ ( t ) = ∑ ν K µν ( t )∆ ν ( t ) .

•

Repeat steps 2-5 until the order parameters converge.

Below we provide a pseudocode algorithm to solve for the propagator elements.

Algorithm 1: Propagator Solver

Data: K x , y , Initial Guesses { Φ glyph[lscript] , G glyph[lscript] } L glyph[lscript] =1 , { A glyph[lscript] , B glyph[lscript] } L -1 glyph[lscript] =1 , Sample count S , Update Speed β Result: Propagator Matrix Σ

1 Solve DMFT equations with Algorithm 2 for order parameters f µ ( t ) , Φ glyph[lscript] µα ( t, s ) , ... ;

2 Draw S samples { u glyph[lscript] µ,n ( t ) } S n =1 ∼ GP (0 , Φ glyph[lscript] -1 ) , { r glyph[lscript] µ,n ( t ) } S n =1 ∼ GP (0 , G glyph[lscript] +1 ) ;

3 Integrate dynamics for each sample to get { h glyph[lscript] µ,n ( t ) , z glyph[lscript] µ,n ( t ) } S n =1 ;

4 Estimate κ functions with Monte Carlo integration, for instance

6 For each sample, compute field sensitivities to error signals, such as ∂h glyph[lscript] µ,n ( t ) ∂ ∆ ( s ) , and kernels

ν ∂h glyph[lscript] µ,n ( t ) ∂ Φ glyph[lscript] αβ ( t ′ ,s ′ ) implicitly using equations (46) (47) ;

7 Use these sensitivities to compute the necessary D tensors such as D Φ glyph[lscript] ∆ µνα = 1 S ∑ n ∈ [ S ] ∂ ∂ ∆ α ( t ′ ) [ φ ( h glyph[lscript] µ,n ( t )) φ ( h glyph[lscript] ν,n ( s )) ] ;

8 Invert U matrix and compute Σ 0 in equation (45);

9 Compute the Schur-complement in equation (43) to handle the response functions ;

The above propagator solver builds on the solution to the DMFT equations which is provided below.

G Leading Correction to the Mean Order Parameters

In this section we use the propagator structure derived in the last section to reason about the leading finite size correction to 〈 q 〉 at width N . Letting the indices i, j, k, n enumerate all entries of the order parameters in q (technically this is a sum over samples and an integral over time for gradient flow), we find the leading Pade Approximant for the mean has the form (App D)

where δ j = √ N ( q j -q ∞ j ) and the derivatives are computed at the saddle point. In the last line, we utilized Wick's theorem and the permutation symmetry of the third derivative ∂ 3 S ∂q i ∂q j ∂q k to evaluate the four point averages in terms of the propagator Σ ij , which was provided in the preceding section E. In practice computing even the full set of second derivatives for the DMFT action to get Σ is quite challenging. Despite the challenge of computing the mean order parameter correction, these corrections are relevant in practice and crucially distinguish the training timescales of deep networks at different widths as we show in Figures 7 and A.4.

G.1 Correction to Mean Predictions and Full MSE Correction

Supposing that we solved for the propagator Σ , using the formalism in the preceeding section, we can compute the O ( N -1 ) correction to the average network prediction error due to finite size. We let 〈 ∆ ( t ) 〉 represent the average of errors over an ensemble of width N networks.

where Σ K ∆ µνν ( t, t ) is the leading covariance (propagator element) between the kernel K µν ( t ) and prediction error ∆ ν ( t ) . We see that the average kernel 〈 K µν ( t ) 〉 (which depends on the finite width N ) plays an important role in characterizing the timescales of the average prediction dynamics. Once this equation is solved for 〈 ∆ µ ( t ) 〉 , the square loss at width N and time t has the form

We will now comment on the structure of the cross term in this above solution. First, if 〈 K 〉 glyph[followsequal] K ∞ and Σ K ∆ is negligible then the average errors at finite width will decay more rapidly than the infinite width model. However, we suspect that in general, 〈 K 〉 -K ∞ contains many negative eigenvalues since signal propagation at finite width tends to reduce the scale of feature kernels [14]. We suspect that this is the cause of the slower dynamics of ensembled predictors for narrower networks in Figure 7 and Figure A.4. Additionally, the term involving Σ K ∆ will generically increase the cross term since the dynamics of ∆ cause its fluctuations to become anti-correlated with the fluctuations in K . In general, it is challenging to make strong definitive statements about the relative scale of these competing effects on the cross term. However, we can say more about this solution in the lazy limit, where we find that the cross term will generically be positive, leading to larger MSE (Appendix H.2).

G.2 Perturbation Theory in Rates rather than Predictions

In experiments on deep CNNs trained on CIFAR-10 in 7 and A.4, we find that the loss curves for the ensemble averaged predictors are effectively time rescaled by a function of network width. In this section, we argue that a proper way to account for this is to compute a perturbation expansion in the exponent which defines the rate of decay of the training errors. To illustrate the point, we first consider the case of a single training example before describing larger datasets. In this case, we consider the change of variables ∆( t ) = e -r ( t ) y . We now treat r as an order parameter of the theory with dynamics

Note that this equation is now a linear relation between two order parameters ( r ( t ) , K ( t ) ), whereas the relation was previously quadratic. In the lazy limit, if K → K -glyph[epsilon1] then r → r -glyph[epsilon1]t , giving an effective rescaling of training time by 1 -glyph[epsilon1] K .

For multiple training examples, we introduce the notion of a transition matrix T ( t ) ∈ R P × P which has dynamics

The solution to the training prediction errors can be obtained at any time t by multiplying the initial condition ∆ (0) = y with the transition matrix ∆ ( t ) = T ( t ) y , where y are the training targets. In this case, the relevant rate matrix , which would be an alternative order parameter is

where log is the matrix logarithm function. Note that in general T ( t ) admits a Peano-Baker series solution [62-64]. In the special case where K ( t ) commutes with ¯ K ( t ) = 1 t ∫ t 0 ds K ( s ) , we obtain the following simplified formula for the rate matrix R

The benefit of this representation is the elimination of coupled order parameter dynamics which are quadratic in fluctuations (in ∆ and K ) into a linear dynamical relation between order parameters R and K . An expansion in R will thus give better predictions at long times t than a direct expansion in ∆ . In the lazy γ → 0 limit, the constancy of K ( t ) = K gives the further simplification R = K t . Working with this representation, we have the following finite width expression for the training loss

where 〈 R 〉 ∼ R ∞ + 1 N R 1 + O ( N -2 ) is the leading correction to the mean R . In this representation, it is clear that finite width can alter the timescale of the dynamics through a correction to the mean of R , as well as contribute an additive correction from fluctuations. This justifies the study perturbation analysis of rates R N as a function of 1 /N in Figures 7 and A.4.

H Variance in the Lazy Limit

We can simplify the propagator equations in the lazy γ → 0 limit. To demonstrate how to use our formalism, we go through the complete process of inverting the Hessian, however, for this case, this procedure is a bit cumbersome. A simplified derivation for the lazy limit can be found below in section H.1 which relies only on linearizing the dynamics around the infinite width solution. In the γ → 0 limit, all of the D tensors vanish and the κ tensors are constant in time. Thus, it suffices to analyze the kernels restricted to t = 0 and study the evolution of the prediction variance ∆ ( t ) .

where { u glyph[lscript] µ } ∼ N (0 , Φ glyph[lscript] -1 ) , { r glyph[lscript] µ } ∼ N (0 , G glyph[lscript] +1 ) . Taking two derivatives with respect to { ˆ Φ glyph[lscript] , ˆ G glyph[lscript] } give terms of the form

Given these we also have the relevant non-vanishing sensitivity tensors

As before we let q 1 = Vec { ∆ µ ( t ) , Φ glyph[lscript] µν , G glyph[lscript] µν , K µν } and q 2 = Vec { ˆ ∆ µ ( t ) , ˆ Φ glyph[lscript] µν , ˆ G glyph[lscript] µν , ˆ K µν } . The propagator has the form

The propagator of interest is Σ q 1 = U -1 [ ∇ 2 q 2 q 2 S ] U -1 glyph[latticetop] . We can exploit the block structure of U to find an inverse

where each sub-block can be computed with the Schur-complement formula. Altogether, we multiply through to get the propagator

Two of these blocks corresponding to K, ∆ are especially important for characterizing the fluctuations of network predictions. The covariance structure for K has the form

Next we use the fact that U -1 ∆Φ = U -1 ∆ K U -1 K Φ and that U -1 ∆ G = U -1 ∆ K U -1 KG , which follows from the block structure of U . Consequently we arrive at the identity

Lastly, we note that, by the Schur-complement formula that U -1 ∆ K = -( I + Θ K ) -1 D ∆ K . Thus, writing ( I + Θ K ) Σ ∆ ( I + Θ K ) glyph[latticetop] = D ∆ K Σ K [ D ∆ K ] glyph[latticetop] as an integral equation, we find

Differentiation with respect to t and s gives a simple differential equation

Let { ψ k } be the eigenvectors of the kernel matrix K . Projecting these dynamics on the eigenspace Σ kglyph[lscript] ( t, s ) = ψ glyph[latticetop] k Σ ( t, s ) ψ glyph[lscript] recovers the equation in the main text

Replacing Σ K = κ recovers the equation (7) in the main text.

H.1 Perturbed Linear System

In this section, we provide a simpler derivation of the lazy limit training error variance dynamics. In this case, we merely perturb the dynamics around its infinite width value ∆ ( t ) = ∆ ∞ ( t ) + glyph[epsilon1] ∆ ( t ) and K = K ∞ + glyph[epsilon1] K , and keep terms only linear in these perturbations. The perturbation glyph[epsilon1] K is fixed in time and the dynamics of glyph[epsilon1] ∆ ( t ) are

Projecting this equation on the eigenspace of K ∞ gives

This immediately recovers the final result of the last section

Qualitatively, the process of computing this linear correction (in glyph[epsilon1] K ) to the dynamics of ∆ is identical to the argument utilized in prior work on perturbative feature learning corrections [11]. In that context, the perturbation is caused by small amounts of feature learning, rather than initialization fluctuations.

H.2 Mean Prediction Error Correction in the Lazy Limit

Using a similar heuristic as in the preceeding section, we now consider the correction to the mean predictor 〈 ∆ µ ( t ) 〉 in the lazy limit. Taylor expanding 〈 ∆ ( t ) 〉 in powers of 1 /N , we find

From the previous section we have that

Projecting these dynamics onto the eigenspace of the kernel gives

where glyph[lscript] = k should be seen as the limit where λ k → λ glyph[lscript] of the above. Thus we find that the leading mean correction to the error solves the following differential equation glyph[negationslash]

We see that at late sufficiently large t , that the terms involving Σ K will dominate. We can gain more intuition by considering the special case of a single training data point where the mean error correction has the form

While the term involving Σ K is positive for all t , K 1 could be positive or negative for a given architecture. If K 1 is positive, then MSE is initially improved at early times but after t > K 1 Σ K the MSE is worse than the infinite width. On the other hand, if K 1 is negative (as we suspect is typically the case), then the MSE will strictly decrease with network width for any time t .

I Two Layer Equations and Time/Time Diagonal

In this section, we analyze two layer networks in greater detail. Unlike the deep network case, two layer networks can be analyzed on the time-time diagonal: ie the dynamics only depend on Φ( t, t ) and G ( t, t ) rather than on all possible off-diagonal pairs of time points. Further, there are no response functions A glyph[lscript] , B glyph[lscript] which complicate the recipe for calculating the propagator (Appendix E).

I.1 A Single Training Point

For a two layer network trained on a single training point with norm constraint | x | 2 = D , we have the following DMFT action

The saddle point equations are

From these equations, we can compute the entries in the Hessian of the DMFT action S . Letting

where κ ( t, s ) = 〈 ( φ ( h ( t )) 2 + g ( t ) 2 )( φ ( h ( s )) 2 + g ( s ) 2 ) 〉 -K ( t ) K ( s ) is the NTK's fourth cumulant. We now vectorize our order parameters over time q = Vec { q ( t ) } t ∈ R + and ˆ q = Vec { ˆ q ( t ) } t ∈ R + and express the full Hessian

The covariance matrix of interest (for q ( t ) ) is thus

where [ Θ K ]( t, s ) = Θ( t -s ) K ( s ) and [ Θ ∆ ]( t, s ) = Θ( t -s )∆( s ) . The above equations allow one to use the infinite width DMFT dynamics for K ( t ) , ∆( t ) to compute the finite size fluctuation dynamics of the kernel K and the error signal ∆ .

I.1.1 Computing Field Sensitivities

In this section, we compute D ( t, s ) by solving for the sensitivity of order parameters. We start with the DMFT field equations

Now, differentiating both sides with respect to ∆( s ′ ) gives

We can compute D Monte carlo by iteratively solving the above equations for each sampled trajectory { h ( t ) , z ( t ) } [65, 46]. Averaging the necessary fields over the Monte Carlo samples will give us the final expressions for D ( t, s ) .

Similarly, the uncoupled kernel variance κ ( t, s ) can be evaluated via Monte Carlo sampling for nonlinear networks.

I.2 Test Point Fluctuation Dynamics

We now are in a position to calculate the test/train kernel and test prediction fluctuations. To do this systematically, we augment S with the test point prediction f glyph[star] and field h glyph[star] and introduce the kernel

K glyph[star] ( t ) = 〈 φ ( h ( t )) φ ( h glyph[star] ( t )) + g ( t ) g glyph[star] ( t ) 〉 . The test prediction f glyph[star] and field h glyph[star] have dynamics

The augmented action for this DMFT has the form

We let q ( t ) = [∆( t ) , f glyph[star] ( t ) , K ( t ) , K glyph[star] ( t )] glyph[latticetop]

Our total covariance matrix / propagator is thus

This is the equation provided in the main text Equation (8).

I.3 Two Layer Linear Network Closed Form

For a linear network on a single data point, we can compute D ( t, s ) and κ ( t, s ) analytically. We start from the field equations

We can make a change of variables v + ( t ) = 1 √ 2 ( h ( t ) + z ( t )) and v -( t ) = 1 √ 2 ( h ( t ) -z ( t )) . We note that v + (0) = 1 √ 2 ( u + r ) and v -(0) = 1 √ 2 ( u -r ) are independent Gaussians. These functions v + ( t ) , v -( t ) satisfy dynamics

Now, we use the fact that v + (0) = 1 √ 2 ( u + r ) and v -(0) = 1 √ 2 ( u -r ) are independent standard normal random variables to compute K ( t ) = 〈 h ( t ) 2 + z ( t ) 2 〉 = 〈 v + ( t ) 2 + v -( t ) 2 〉

This operator is causal ( D ( t, s ) = 0 for s > t ) as expected and vanishes as t → 0 . If we take γ → 0 , we have D ( t, s ) → 0 which agrees with our reasoning that fields h, z only depend on ∆ in the feature learning regime. Since all fields are Gaussian in the linear network case, we can use Wick's theorem to obtain the exact uncoupled kernel variance in the two layer case.

The v ± ( t ) functions are those given above. Using the fact that 〈 v + (0) 2 〉 = 〈 v -(0) 2 〉 = 1 allows us to easily compute the single site average above.

J Multiple Samples with Whitened Data

In this section, we analyze the role that sample number plays in dynamics in a simplified model of a two layer linear network trained on whitened data. Concretely, we assume that x µ · x ν D = δ µν . The field equations for preactivations h µ ( t ) and pregradients z ( t ) obey

We will assume the targets have unit norm | y | 2 = 1 and we define the projection of ∆ onto the target as ∆ y ( t ) = y · ∆ ( t ) . The other P -1 orthogonal components are denoted ∆ ⊥ ( t ) so that ∆ = ∆ y ( t ) y + ∆ ⊥ ( t ) with ∆ ⊥ ( t ) · y = 0 . At infinite width, ∆ ⊥ = 0 and our field equations become

However, at finite width N , the off-target predictions ∆ ⊥ fluctuate over random initialization. To model all of the fluctuations simultaneously, we consider the following action

which enforces the constraint that ∆ µ ( t ) = y µ -1 γ 〈 z ( t ) h µ ( t ) 〉 at infinite width. The Hessian over order parameters q = Vec { ∆ µ ( t ) , ˆ ∆ µ ( t ) } has the form

We thus get the following covariance for predictions Σ ∆ = ( γ I + D ) -1 κ [ ( γ I + D ) -1 ] glyph[latticetop] . We now compute the necessary components of the D tensor

In the last line, we used the fact that these equations are to be evaluated at the mean field infinite width stochastic process where ∆ ⊥ ( t ) = 0 . To compute the sensitivity tensor D , we find the following equations for our correlators of interest:

glyph[negationslash]

glyph[negationslash]

We therefore see that the components of D decouple over indices. In the y direction, we have the following equations

where the correlators must be solved self-consistently. We will provide this solution in one moment, but first, we will look at the orthogonal directions. For the P -1 orthogonal directions, we obtain the explicit formula for D in each of these directions

Now, we return to D y . To solve these equations we utilize the change of variables employed in the single sample case v + ( t ) = 1 √ 2 ( h y ( t ) + z ( t )) , v -( t ) = 1 √ 2 ( h y ( t ) -z ( t )) (see Appendix I.3). This orthogonal transformation decouples the dynamics

As a consequence, the field derivatives close

The correlator of interest is

So we get that

Similarly, we can derive the on-target and off-target uncoupled variances κ y ( t, s ) and κ ⊥ ( t, s ) , which satisfy

Using these functions, we arrive at the following variance for each of the P dimensions

Using the fact that all ∆ ⊥ variables are independent and identically distributed under the leading order picture, the expected training loss has the form

where 〈 ∆ y -∆ ∞ y 〉 = 1 N ∆ 1 y ( t ) + O ( N -2 ) . We note that the bias correction if O ( N -1 ) while the variance is O ( P/N ) . We compare the above leading order theory with and without the bias correction in Appendix Figure A.2.

K Online Learning

Our technology for computing finite size effects can easily be translated to a setting where the neural network is trained in an online fashion, disregarding the effect of SGD noise. At each step, we compute the gradient over the full data distribution p ( x ) . Focusing on MSE loss, we study the following equation

where K ( x , x ′ ; t ) is the dynamic NTK and ∆( x , t ) = y ( x ) -f ( x , t ) is the prediction error. In general the distribution involves integration over an uncountable set of possible inputs x . To remedy this, we utilize a countable orthonormal basis of functions for the data distribution { ψ k ( x ) } ∞ k =1 . For example, if p ( x ) were the isotropic Gaussian density for N (0 , I ) , then ψ k could be Hermite polynomials. We expand ∆ and K in this basis ψ k , and arrive at the following differential equation

By orthonormality, the average turned into a sum over all possible orthonormal functions { ψ k } . We note that since K is evolving in time, there is not generally a fixed basis of functions that diagonalize K , resulting in the couplings across eigenmodes in Equation (111). Since, in online learning, there is no distinction between the training and test distribution, our error of interest is simply L ( t ) = ∑ k ∆ k ( t ) 2 . To obtain the finite size corrections to this quantity, we compute the joint propagator for all variables { K kglyph[lscript] ( t ) , ∆ k ( t ) } . If we wanted to pursue a perturbation theory in rates (Appendix G.2), we could again define a transition matrix T and rate matrix R ( t ) as

We can then obtain ∆ = exp( -R ( t )) y , where y k = E x ψ k ( x ) y ( x ) . Since R has a finite size mean correction and finite size fluctuations, so too does the error ∆ k ( t ) and the loss L (Appendix G.2).

K.1 Two Layer Networks

In the two layer case, instead of tracking kernels, we could instead deal with the distribution over read-in vectors w ∈ R D and readout scalars a ∈ R as in the original works on mean field networks [6, 66]. When training on the population risk equations for x ∼ N (0 , I )

The action has the form

The Hessian over q = { ∆ µ ( t ) , ˆ ∆ µ ( t ) } is

The above equations could be solved and then used to compute D ∆ ( t, x ; s, x ′ ) which must then be inverted to get the observed prediction variance.

K.2 Linear Activations

Using the ideas in the preceding sections, we can make more progress in the case of a two layer linear network in the online learning setting. The key idea is to track the kernel and prediction error projections onto the space of linear functions. In this case we get the following DMFT over the order parameter β ( t ) = 1 N W glyph[latticetop] a ∈ R D .

At infinite width, we see that the dynamics can be reduced to tracking the projection of the weights w and β on the β glyph[star] direction. The D -1 off-target dimensions vanish β ⊥ ( t ) = 0 . At infinite width, we arrive at the alignment dynamics studied in prior work [64, 9]

We note that β ( t ) = β ( t ) β glyph[star] and that M has only one special eigenvector β glyph[star] with eigenvalue m glyph[star] ( t ) . It thus suffices to track evolution in this single direction

We note that this equation is identical to the differential equation for a single training example in Appendix J. Here β glyph[star] -β ( t ) plays the role of ∆ y ( t ) and m glyph[star] ( t ) plays the role of the kernel K y ( t ) . A key observation is the conservation law 4 γ 2 d dt β ( t ) 2 = d dt m glyph[star] ( t ) 2 , from which it follows that m glyph[star] ( t ) 2 -4 = 4 γ 2 β ( t ) [9]

This is identical to the differential equations for a single sample (producing prediction f ( t ) and kernel K ( t ) ) if the following substitutions are made

We now proceed to compute finite size corrections starting from the action

The necessary ingredients are

Similarly we have to compute the sensitivity tensor

We start from the dynamics

Next, we have to calculate causal derivatives for fields

Following an identical argument as in J, we see that D has block diagonal structure with D β glyph[star] ( t, s ) on the β glyph[star] β glyph[latticetop] glyph[star] direction and D ⊥ ( t, s ) in any of the D -1 remaining directions

Similarly, κ ( t, s ) has a similar decomposition

The processes have the following equations at infinite width

As a consequence we note that 〈 w ⊥ ( t ) a ( s ) 〉 = 0 so that κ ⊥ ( t, s ) = 〈 a ( t ) a ( s ) 〉 . Letting v + ( t ) = 1 √ 2 ( w β glyph[star] ( t )+ a ( t )) and v -( t ) = 1 √ 2 ( w β glyph[star] ( t )+ a ( t )) , we find the same decoupled stochastic processes as in Appendix I.3.

We can use these equations to perform the necessary averages for κ β glyph[star] and D β glyph[star] . Lastly, we use

to evaluate D ⊥ ( t, s ) . The observed covariances are just

We note that these expressions are identical to those in Appendix J under the substitution β glyph[star] -β ( t ) → ∆( t ) and D → P . Thus the expected test risk is

This recovers the variance we obtained in the multiple-sample whitened data case J.

K.3 Connections to Offline Learning in Linear Model

As in the offline case, in Fig. 4 (c) and (d) we see that the variance contribution to test loss | β -β glyph[star] | 2 increases with input dimension D . We note that this perturbative effect to the loss dynamics is reminiscent of the deviations from mean field behavior studied in SGD [43, 44], though this present work concerns fluctuations driven by initialization variance rather than stochastic sampling of data. In Fig. 4 (e) we show that richer networks have lower variance at fixed N . Similarly, leading order theory for richer networks more accurately captures their dynamics as D/N increases (Fig. 4 (f)).

L Deep Linear Networks

For deep linear networks, the fields h glyph[lscript] µ ( t ) , g glyph[lscript] µ ( t ) are Gaussian and have the following self-consistent equations

B glyph[lscript] µν ( t, s ) = 〈 ∂h glyph[lscript] µ ( t ) ∂r ν ( s ) 〉 [9]. Therefore, we express the action as a differentiable function of the order parameters by integrating over the Gaussian field distribution. For concreteness, we vectorize our fields over time and samples h glyph[lscript] = Vec { h glyph[lscript] µ ( t ) } { µ ∈ [ P ] ,t ∈ R + } , g glyph[lscript] = Vec { g glyph[lscript] µ ( t ) } { µ ∈ [ P ] ,t ∈ R + } we consider the contribution of a single hidden layer.

where C glyph[lscript] µν ( t, s ) = γ Θ( t -s ) [ A glyph[lscript] -1 µν ( t, s ) + H glyph[lscript] -1 µν ( t, s )∆ ν ( s ) ] and D glyph[lscript] µν ( t, s ) = γ Θ( t -s ) [ B glyph[lscript] µν ( t, s ) + G glyph[lscript] +1 µν ( t, s )∆ ν ( s ) ] . Performing the joint Gaussian integrals over ( h glyph[lscript] , g glyph[lscript] , ˆ h glyph[lscript] , ˆ g glyph[lscript] ) we find

We can then automatically differentiate the DMFT action to get the propagator. For example, for a three layer linear network, the full DMFT action has the form

where C 1 = γ Θ ∆ and C 2 = γ Θ ∆ glyph[circledot] H 1 + γ A and D 1 = γ Θ ∆ glyph[circledot] G 2 + γ B and D 2 = γ Θ ∆ . This above example can be extended to deeper networks. The total size of the block matrices which we compute determinants over is 4 PT × 4 PT for a dataset of size P trained for T steps.

M Discrete Time Dynamics and Edge of Stability Effects

Large step size effects can induce qualitatively different dynamics in neural network training. For instance, if the step size exceeds that required for linear stability with the initial kernel, the kernel can decrease in order to stabilize the dynamics [57]. Alternatively, during training the kernel may exhibit a 'progressive sharpening" phase where its top eigenvalue grows before reaching a stability bound set by the learning rate [19]. It is therefore well motivated to study how dynamics in this regime alter finite size effects in neural networks. We will first solve a special model which was considered in prior work [57]: a two layer linear network trained on a single training point. We will then provide the full DMFT equations for the discrete time case and provide an outline for how one could obtain finite size effects in that picture.

M.1 Two Layer Linear Equations

In a two layer linear network, the DMFT equations are

The NTK has the form K ( t ) = 〈 h ( t ) 2 + z ( t ) 2 〉 . We can easily show that the kernel and error have coupled dynamics

These equations define the infinite width evolution of ∆( t ) and K ( t ) . Already at this level of analysis, we can reason about the evolution of K ( t ) . In the small η limit, we could disregard terms of order O ( η 2 ) and arrive at the following gradient flow approximation for K ( t ) ∼ 2 √ 1 + γ 2 f ( t ) 2 [9]. This evolution will not reach the edge of stability provided that η < 1 √ 1+ γ 2 y 2 . For large γ and y = 1 , this leads to the constraint ηγ < 1 . However, if η exceeds this bound, the gradient flow approximation is no longer reasonable and the system reaches an edge of stability effect as shown in Figure 6.

To calculate the finite size effects, we need to compute κ and D ( t, s ) = ∂ ∂ ∆( s ) 〈 h ( t ) 2 + z ( t ) 2 〉 . To evaluate these quantities we utilize the same change of variables employed in Appendix I.3. In discrete time, these decoupled equations are

Given ∆( t ) , these can be expressed as linear systems of equations. Now, we can easily compute the uncoupled kernel variance

Similarly, we can calculate D ( t, s ) by using the fact 〈 h ( t ) 2 + z ( t ) 2 〉 = 〈 v + ( t ) 2 + v -( t ) 2 〉

These can be directly solved as a linear system of equations.

N Computing Details

Experiments for Figures 3, 6 and 2 were conducted on a Google Colab GPU with JAX. Experiments for Figures 5, A.3, 7 were performed on a NVIDIA SMX4-A100-80GB GPU. The total compute required for all Figures in the paper took around 4 hours. Jupyter Notebooks to reproduce plots can be found at https://github.com/Pehlevan-Group/dmft_fluctuations .

Peer reviews:
comments: SUMMARY:
The authors study a finite width correct to the Dynamical Mean Field Theory (DMFT) of finite depth neural networks in the feature learning regime. While I will be the first to admit that I am not an expert on the DMFT calculations, the authors did produce very convincing simulations capturing interesting properties of finite size networks. In particular, the edge of stability behaviour at large learning rates seems to be modelled by the finite width correction. 

I hope the authors can help me clarify some questions regarding the implementation of the solver, after which I would be happy to raise my score to accept.

STRENGTHS:
The authors develop a strong theory of finite width neural networks capable of making accurate predictions.

WEAKNESSES:
N/A

QUESTIONS:
I would like the authors to provide some intuitions about how the DMFT equations are being solved numerically. I understand this may be the contents of a previous paper (Bordelon and Pehlevan, 2022), but I am having trouble understanding this part. 

1. Based on my understanding, all the elements of $q$ are processes index by time $t$. So what does it mean to compute elements of $q$ or $\Sigma$? Are the authors time discretizing these processes? 
2. How should I understand the role of the propagator $\Sigma$, and how the authors solve for this object first independently of the elements of $q$? 
3. How should I understand the saddle point solver in Algorithm 2? Here I am just asking for what is happening within the algorithm, as I find reading the pseudocode quite difficult to understand. 

I would be happy to continue this discussion during the rebuttal period, and once again I would raise my score once I find these questions addressed.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBuilding on past work which set up a DMFT (Dynamical Mean Field Theory) for fully connected networks in the infinite width limit (where the width of each layer tends to infinity), this paper reasons about the *fluctuations* around the infinite width limit. This is important because for finite sized neural networks, these fluctuations are large enough that they are an important part of the network training dynamics. The description of the DMFT is theoretically complicated and cannot be solved exactly, but they enable simulations which confirm that the DMFT captures the behaviour of finite sized real networks quite accurately. As an application, this DMFT theory is used to understand bias, training rates, and variance in realistic tasks and the special case of 2-layer networks (where the theory is quite a bit simpler) is investigated in more detail. This theory has the potential to open the door to many potential future uses that explain how neural networks learn.\n\nSTRENGTHS:\nThis paper lays out a framework for theoretical understanding of deep neural networks that incorporates the effect of finite width. This is a  problem that has received a lot of attention since the original NTK/infinite width limits came out, and as far as I know, the approach here is novel and powerful. The paper is well written for the most part although familiarity with the infinite width DMFTs is assumed. \n\nOverall, the fact that the theory and the simulations agree very well is quite impressive, and I think the ideas in the paper are quite ambitious because they can be used for almost any kind of question one might have about the theoretical evolution of the DNN. This paper has the potential to be the basis for future work which uses the theory developed here to investigate questions about how DNNs learn.\n\nWEAKNESSES:\nThe main weakness of the paper is that its a bit spread thin at times: both the theory and a few different applications are covered, but it seems like the authors were trying to make it all fit and I would have liked more detail in a few spots. This is largely due to the page limit of the submission. I personally would have found it to be a stronger paper if a single really clear example was presented in a lot of detail. (Although again, I completely understand that this is largely pressure from the conference format to try and do a lot of stuff)  \n\nThe other main \"weakness\" of the paper (which is strictly speaking a limitation of the audience of the paper) is that to understand it, you need to be familiar with the previous DMFT on which this paper is built. The authors include a very short section called \"Review of Dynamical Mean Field Theory\" citing [9],[46] as a review, but this section is extremely sparse for actually understanding what is going on. I essentially had to read [9] in its entirety first to understand what was going on in this paper. (Also the reference [46] could not be found since only authors and title are given...where would one find this reference?) In my view, this weakness could be mitigated by just being more honest with the reader up front about this...for example [9],[46] should be cited at point 1 in the list of contributions to make it more clear the dependence and what is/isn't actually explained in this paper. \n\nAnother (related) \"weakness\" is that the paper relies quite heavily on physics technology and jargon to reach its conclusions. The fact that the results are so heavily entrenched in physics jargon like \"order parameters\" or \"propagator\" makes this paper less likely to have a broad impact on the deep learning community. The authors would add a lot of value to the work by attempting to make a \"translation guide\" to help people who don't have the same physics background understand what is going on in more detail.\n\nQUESTIONS:\n* Suggestion: Eqn (3) seems like a very important main result: a bit more discussion on the next order term and why its actually size 1/N (even though it starts with an N) would be helpful. \n* I am wondering if you had any numerical simulations where you checked the rate of covergence between the theory here and simulations (which should presumably be like 1/N^2?). Something like Figure 2a but comparing to the fluctuation predictions rather than to $q_\\infty$.\n* One high-level question I had: It seems that if the learning rate is fixed and not scaled, the the fluctuations and the effect of the learning rate are both on the same scale $1/\\sqrt{N}$. 1. Is that correct? and 2. Does your theory work to analyze what's going on in that case?\n\nHere is a list of other minor errors/suggestions I found while reading:]\n* Eqn (2): use definition equal with three lines to be consistent with definitions later on\n* Line 179 vs 185: Is there a difference between $K$ and $K_\\infty$? If so what is it?\n* Line 188: What does the subscript 0 in \"$Cov_0$\" mean here?\n*  Section 6.1: I was able to more-or-less piece this together, but I think it would be a lot more understandable if you gave explit definitions for $K(t)$ and $K_\\ast(t)$\n* Line 232-235: I think it would be a lot clearer to write out the definitions here of the two new $\\Delta$'s in terms of $\\Delta_\\mu$. \n* References [15],[21],[33], [46], [53] has only author/title but not where published/where to find.\n* Reference [22] missing a title?\n* I would also check the arxiv only referenes e.g. [11],[13],[20],[27],[34],[57],[63] to make sure there isnt a conference or journal version that is now published.\n\nLIMITATIONS:\nA potential criticism is the physics level of rigor used in Appendix C which is used to establish the main results. The manipulations carried out in the proof of Appendix C certainly seem plausible, and I believe the community as a whole is ok with this level of rigour, but the authors could be a bit more clear about what they mean by \"proof\" in the main paper. It is not a mathematically rigour proof (which would involve all sorts of techincal assumptions), but rather a physics-type statement that holds assuming the usual expansions can be carried out without obstructions. To reiterate: I think the actual work is fine, but they could be a bit more honest about how it is \"proven\" and the level of rigour in the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the problem of analytical description of the rich (feature learning) dynamics of neural networks. To achieve this, the authors use previously introduced dynamical mean field theory (DMFT), which identifies several key characteristics of the problem - order parameters - defines their probability distribution throughout the training using a path integral representation. The paper considers $\\mu P$ parametrization of the network, which is known to display feature learning behavior at infinite width (in contrast to popular NTK and Standard parametrizations) with feature learning strength controlled by parameter $\\gamma$.\n\nIn the paper, the authors focus on leading width $O(N^{-1})$ corrections to the infinite width limit. Technically, this is achieved by taking into account quadratic (Gaussian) fluctuation around saddle point $\\mathbf{q}_\\infty$ of DMFT action $S(\\mathbf{q})$. After deriving general equations governing $O(N^{-1})$ corrections, the authors consider a number of simplified scenarios where these equations can be solved analytically. Also, the authors validate qualitative conclusions of their theory in the non-synthetic experiments with CNN trained on CIFAR10.\n\nSTRENGTHS:\nThe approach chosen by the authors - identifying saddle point $\\mathbf{q}_\\infty$ of the system's action $S(\\mathbf{q})$ and investigating it together with Gaussian fluctuation around $\\mathbf{q}_\\infty$ - is a workhorse of various branches of theoretical physics where behavior of a complex system needs to be analyzed. In physics, this approach not only provides a very sizable portion of available analytical results but, for example, also provides SOTA results for numerical modeling of realistic strongly-correlated materials [Kotliar 2004](https://pubs.aip.org/physicstoday/article-abstract/57/3/53/755526/Strongly-Correlated-Materials-Insights-From?redirectedFrom=fulltext), [Vollhardt 2019](https://arxiv.org/abs/1910.12650). Thus, realizing this general strategy is one of the fundamental directions within deep learning theory. Also, as noted by the authors, the DMFT approach is non-perturbative w.r.t. feature learning strength  - a feature that is mostly absent in other approaches to NN dynamics away from kernel regime.\n\nWEAKNESSES:\n* In many cases, solving DMFT equations analytically seems to be intractable. This significantly renders the main purpose of the proposed theory - obtaining analytical insights into the network dynamics.\n* The DMFT equations are bulky, which could make working with them quite exhaustive.\n* I believe it is hard to understand main DMFT ingredients - order parameters $\\mathbf{q}$ and their action $S(\\mathbf{q})$ - from the current paper alone. Most probably, a careful reading of the original DMFT paper [9] is required to understand this paper. However, this is not the author's fault but rather a consequence of the chosen approach.\n\nQUESTIONS:\n* What dataset size $P$ were used for main CIFAR10 experiment of fig. 10? Due to mentioned $O(P^4T^4)$ memory requirements of storing the propagators, it seems to be impossible to work with the full ($P=50000$) CIFAR10 size. Also, what is the computational complexity of solving DMFT equations? Both for saddle point $\\mathbf{q}_\\infty$ and $O(N^{-1})$ corrections if they are different.\n* While in most of the presented plots DMFT accurately describes the experiments, for the big sample sizes (fig. 3) there is a significant disagreement between theory and experiment. Do you think this a fundamental limitation DMFT on the level of leading order correction (ignoring $O(N^{-2})$ and beyond)? Or maybe it is because in the experiment you considered whitened data, whereas in the more realistic scenarios, the data typically has a low effective dimension (e.g. measured by the decay of data covariance eigenvalues)?\n\nLIMITATIONS:\nThe authors discuss the limitations of their current approach. The limitations mainly come from 1) the inability to numerically solve DMFT equations for large-scale problems (e.g. due to $O(P^4 T^4)$ propagator size) and  2) the need to consider higher order expansion around the saddle point $\\mathbf{q}_\\infty$.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}