SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Blake Bordelon & Cengiz Pehlevan
John Paulson School of Engineering and Applied Sciences,
Center for Brain Science,
Kempner Institute for the Study of Natural & Artificial Intelligence,
Cambridge MA, 02138
blake_bordelon@g.harvard.edu, cpehlevan@g.harvard.edu
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{\mathrm{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.
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  $\mu$ -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 finite-size 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  $\mathcal{O}(1 / \mathrm{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:
1. 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  $\mathcal{O}(1 / \mathrm{width})$  variance of feature and gradient kernels throughout training.
2. 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.
3. 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.
4. 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.
5. 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.
6. 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.
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  $\mu 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  $\mathcal{O}(1 / \sqrt{\mathrm{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  $\Phi^{\ell}$  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.
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  $\mathbf{x}_{\mu} \in \mathbb{R}^{D}$  to hidden preactivations  $\pmb{h}_{\mu}^{\ell} \in \mathbb{R}^{N}$  in layers  $\ell \in \{1, \dots, L\}$  and finally output  $f_{\mu}$
$$
f _ {\mu} = \frac {1}{\gamma N} \boldsymbol {w} ^ {L} \cdot \phi \left(\boldsymbol {h} _ {\mu} ^ {L}\right), \quad \boldsymbol {h} _ {\mu} ^ {\ell + 1} = \frac {1}{\sqrt {N}} \boldsymbol {W} ^ {\ell} \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell}\right), \quad \boldsymbol {h} _ {\mu} ^ {1} = \frac {1}{\sqrt {D}} \boldsymbol {W} ^ {0} \boldsymbol {x} _ {\mu}, \tag {1}
$$
where  $\gamma$  is a scale factor that controls feature learning strength, with large  $\gamma$  leading to rich feature learning dynamics and the limit of small  $\gamma \rightarrow 0$  (or generally if  $\gamma$  scales as  $N^{-\alpha}$  for  $\alpha > 0$  as  $N \to \infty$ , NTK parameterization corresponds to  $\alpha = \frac{1}{2}$ ) gives lazy learning where no features are learned [4, 7, 9]. The parameters  $\pmb{\theta} = \{\pmb{W}^0, \pmb{W}^1, \dots, \pmb{w}^L\}$  are optimized with gradient descent or gradient flow  $\frac{d}{dt} \pmb{\theta} = -N \gamma^2 \nabla_{\pmb{\theta}} \mathcal{L}$  where  $\mathcal{L} = \mathbb{E}_{\pmb{x}_{\mu} \in \mathcal{D}} \ell(f(\pmb{x}_{\mu}, \pmb{\theta}), y_{\mu})$  is a loss computed over dataset  $\mathcal{D} = \{(x_1, y_1), (x_2, y_2), \dots, (x_P, y_P)\}$ . This parameterization and learning rate scaling ensures that  $\frac{d}{dt} f_{\mu} \sim \mathcal{O}_{N,\gamma}(1)$  and  $\frac{d}{dt} h_{\mu}^{\ell} = \mathcal{O}_{N,\gamma}(\gamma)$  at initialization. This is equivalent to maximal update parameterization ( $\mu$ 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)  $\mathbf{q} = \mathrm{Vec}\{f_{\mu}(t),\Phi_{\mu \nu}^{\ell}(t,s),G_{\mu \nu}^{\ell}(t,s),\ldots \}$  [9], which include feature and gradient kernels [9, 54]
$$
\Phi_ {\mu \nu} ^ {\ell} (t, s) \equiv \frac {1}{N} \phi \left(\boldsymbol {h} _ {\mu} ^ {\ell} (t)\right) \cdot \phi \left(\boldsymbol {h} _ {\nu} ^ {\ell} (s)\right), \quad G _ {\mu \nu} ^ {\ell} (t, s) \equiv \frac {1}{N} \boldsymbol {g} _ {\mu} ^ {\ell} (t) \cdot \boldsymbol {g} _ {\nu} ^ {\ell} (s), \tag {2}
$$
where  $g_{\mu}^{\ell}(t) = \gamma N\frac{\partial f_{\mu}(t)}{\partial h_{\mu}^{\ell}(t)}$  are the back-propagated gradient signals. Further, for width-  $N$  networks the distribution of these dynamical variables across weight initializations (from a Gaussian distribution  $\pmb{\theta} \sim \mathcal{N}(0,\mathbf{I})$ ) is given by  $p(\pmb{q}) \propto \exp(NS(\pmb{q}))$ , where the action  $S(\pmb{q})$  contains interactions between neuron activations and the kernels at each layer [9].
The DMFT introduced in [9] arises in the  $N\to \infty$  limit when  $p(\boldsymbol {q})$  is strongly peaked around the saddle point  $q_{\infty}$  where  $\frac{\partial S}{\partial q} |_{q_{\infty}} = 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_\mu^\ell (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_\mu^\ell (t)$ . Kernels in (2) are now computed as averages over these infinite-width single site processes  $\Phi_{\mu \nu}^{\ell}(t,s) = \left\langle \phi (h_{\mu}^{\ell}(t))\phi (h_{\nu}^{\ell}(s))\right\rangle ,G_{\mu \nu}^{\ell}(t,s) = \left\langle g_{\mu}^{\ell}(t)g_{\nu}^{\ell}(s)\right\rangle$  where averages arise from the  $N\rightarrow \infty$  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.
Order params. q	Action S(q)	Propagator Σ	Single Site Density
Concentrating variables	q's log-density	Asymptotic Covariance	Neuron Marginals
Table 1: Relationship between the physics and ML terminology for the central objects in this paper. The  $q$  which concentrate at infinite width, but fluctuate at finite width  $N$ . This paper is primarily interested in  $\Sigma$ , the asymptotic covariance of the order parameters.
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  $\pmb{q}$  fluctuate in a  $\mathcal{O}(1 / \sqrt{N})$  neighborhood of  $q_{\infty}$  [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(\mathbf{q})$  across initializations, which we denote by  $\langle O(\mathbf{q})\rangle_N$ , admits an expansion of the form whose leading terms are
$$
\langle O (\boldsymbol {q}) \rangle_ {N} = \frac {\int d \boldsymbol {q} \exp (N S [ \boldsymbol {q} ]) O (\boldsymbol {q})}{\int d \boldsymbol {q} \exp (N S [ \boldsymbol {q} ])} = \langle O (\boldsymbol {q}) \rangle_ {\infty} + N [ \langle V (\boldsymbol {q}) O (\boldsymbol {q}) \rangle_ {\infty} - \langle V (\boldsymbol {q}) \rangle_ {\infty} \langle O (\boldsymbol {q}) \rangle_ {\infty} ] + \dots , \tag {3}
$$
where  $\langle \rangle_{\infty}$  denotes an average over the Gaussian distribution  $\pmb{q} \sim \mathcal{N}\left(\pmb{q}_{\infty}, -\frac{1}{N}\left(\nabla_{\pmb{q}}^{2}S[\pmb{q}_{\infty}]\right)^{-1}\right)$  and the function  $V(\pmb{q}) \equiv S(\pmb{q}) - S(\pmb{q}_{\infty}) - \frac{1}{2} (\pmb{q} - \pmb{q}_{\infty})^{\top}\nabla_{\pmb{q}}^{2}S(\pmb{q}_{\infty})(\pmb{q} - \pmb{q}_{\infty})$  contains cubic and higher terms in the Taylor expansion of  $S$  around  $\pmb{q}_{\infty}$ . 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  $\mathcal{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):  $\boldsymbol{\Sigma} = -\left[\nabla^{2}S(\boldsymbol{q}_{\infty})\right]^{-1}$ . This object can be shown to capture leading order fluctuation statistics  $\left\langle (\boldsymbol{q} - \boldsymbol{q}_{\infty})(\boldsymbol{q} - \boldsymbol{q}_{\infty})^{\top}\right\rangle_{N} = \frac{1}{N}\boldsymbol{\Sigma} + \mathcal{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  $\gamma$  [7, 17, 18]. To calculate such corrections, in App. E, we provide a complete description of Hessian  $\nabla_q^2 S(\boldsymbol{q})$  and its inverse (the propagator) for a depth- $L$  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  $\kappa$  which describe the uncoupled variances of the kernels, such as
$$
\kappa_ {\mu \nu \alpha \beta} ^ {\Phi^ {\ell}} (t, s, t ^ {\prime}, s ^ {\prime}) \equiv \left\langle \phi \left(h _ {\mu} ^ {\ell} (t)\right) \phi \left(h _ {\nu} ^ {\ell} (s)\right) \phi \left(h _ {\alpha} ^ {\ell} \left(t ^ {\prime}\right)\right) \phi \left(h _ {\beta} ^ {\ell} \left(s ^ {\prime}\right)\right) \right\rangle - \Phi_ {\mu \nu} ^ {\ell} (t, s) \Phi_ {\alpha \beta} ^ {\ell} \left(t ^ {\prime}, s ^ {\prime}\right) \tag {4}
$$
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 perturbations of order parameters, such as
$$
D _ {\mu \nu \alpha \beta} ^ {\Phi^ {\ell} \Phi^ {\ell - 1}} (t, s, t ^ {\prime}, s ^ {\prime}) \equiv \frac {\partial \left\langle \phi \left(h _ {\mu} ^ {\ell} (t)\right) \phi \left(h _ {\nu} ^ {\ell} (s)\right) \right\rangle}{\partial \Phi_ {\alpha \beta} ^ {\ell - 1} \left(t ^ {\prime} , s ^ {\prime}\right)}, D _ {\mu \nu \alpha} ^ {G ^ {\ell} \Delta} (t, s, t ^ {\prime}) \equiv \frac {\partial \left\langle g _ {\mu} ^ {\ell} (t) g _ {\nu} ^ {\ell} (s) \right\rangle}{\partial \Delta_ {\alpha} (t ^ {\prime})}, \tag {5}
$$
where  $\Delta_{\mu}(t) = -\frac{\partial\ell(f_{\mu},y_{\mu})}{\partial f_{\mu}} |_{f_{\mu}(t)}$  are error signal for each data point.
Abstractly, we can consider the uncoupled variances  $\kappa$  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  $\Phi^{\ell}$  were perturbed,  $D^{\Phi^{\ell +1},\Phi^{\ell}}$  would quantify the resulting perturbation to  $\Phi^{\ell +1}$  through the fields  $h^{\ell +1}$ .
In App. E, we calculate  $\kappa$  and  $D$  tensors, and show how to use them to calculate the propagator. As an example of our results:
Proposition 2 Partition  $\mathbf{q}$  into primal  $\mathbf{q}_1$ $= \text{Vec}\{f_{\mu}(t), \Phi_{\mu \nu}^{\ell}(t, s)...\}$  and conjugate variables  $\mathbf{q}_2$ $= \text{Vec}\{\hat{f}_{\mu}(t), \hat{\Phi}_{\mu \nu}^{\ell}(t, s)...\}$ . Let  $\kappa = \frac{\partial^2}{\partial q_2 \partial q_2^\top} S[q_1, q_2]$  and
$\pmb{D} = \frac{\partial^2}{\partial q_2 \partial q_1^\top} S[q_1, q_2]$ , then the propagator for  $q_1$  has the form  $\pmb{\Sigma}_{q_1} = \pmb{D}^{-1} \pmb{\kappa} \left[ \pmb{D}^{-1} \right]^\top$  (App E). The variables  $q_1$  are
related to network observables, while conjugates  $\mathbf{q}_2$  arise as Lagrange multipliers in the DMFT calculation. From the propagator  $\Sigma_{\mathbf{q}_1}$  we can read off the variance of network observables such as  $NVar(f_{\mu}) \sim \Sigma_{f_{\mu}}$ .
Figure 1: The directed causal graph between DMFT order parameters (blue) and fields (green) defines the  $D$  tensors of our theory. Each arrow represents a causal dependence.  $K$  denotes the NTK.
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  $\Sigma$  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  $\gamma \rightarrow 0$  where features are frozen. To prevent divergence of the network in this limit, we use a background subtracted function  $\tilde{f}(\boldsymbol{x},\boldsymbol{\theta}) = f(\boldsymbol{x},\boldsymbol{\theta}) - f(\boldsymbol{x},\boldsymbol{\theta}_0)$  which is identically zero at initialization [4]. For mean square error, the  $N\to \infty$  and  $\gamma \rightarrow 0$  limit is governed by  $\frac{\partial\tilde{f}(\boldsymbol{x})}{\partial t} = \mathbb{E}_{\boldsymbol{x}'\sim \mathcal{D}}\Delta (\boldsymbol{x}')K(\boldsymbol{x},\boldsymbol{x}')$  with  $\Delta (\pmb {x}) = y(\pmb {x}) - \tilde{f} (\pmb {x})$  (for MSE) and  $K$  is the static (finite width and random) NTK. The finite-  $N$  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  $\gamma \rightarrow 0$  limit, it has constant initialization variance through training. Further, all sensitivity blocks of the Hessian involving the kernels and the prediction errors  $\Delta$  (such as the  $D^{\Phi^{\ell},\Delta}$ ) vanish. We represent the covariance of the NTK as  $\kappa (\pmb {x}_1,\pmb {x}_2,\pmb {x}_2,\pmb {x}_3) = NCov(K(\pmb {x}_1,\pmb {x}_2),K(\pmb {x}_3,\pmb {x}_4))$ . To identify the dynamics of the error  $\Delta$  covariance, we relate  $K$ , the finite width NTK, to  $K_{\infty}$  which is the deterministic infinite width NTK  $K_{\infty}$ . We consider the eigendecomposition of the infinite-width NTK  $K_{\infty}(\pmb {x},\pmb{x}^{\prime}) = \sum_{k}\lambda_{k}\psi_{k}(\pmb {x})\psi_{k}(\pmb{x}^{\prime})$  with respect to the training distribution  $\mathcal{D}$ , and decompose  $\kappa$  in this basis.
$$
\kappa_ {k \ell m n} = \left\langle \kappa \left(\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2}, \boldsymbol {x} _ {3}, \boldsymbol {x} _ {4}\right) \psi_ {k} \left(\boldsymbol {x} _ {1}\right) \psi_ {\ell} \left(\boldsymbol {x} _ {2}\right) \psi_ {n} \left(\boldsymbol {x} _ {3}\right) \psi_ {m} \left(\boldsymbol {x} _ {4}\right) \right\rangle_ {\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2}, \boldsymbol {x} _ {3}, \boldsymbol {x} _ {4} \sim \mathcal {D}}, \tag {6}
$$
where averages are computed over the training distribution  $\mathcal{D}$ .
Proposition 3 For MSE loss, the prediction error covariance  $\Sigma^{\Delta}(t,s) = N\mathrm{Cov}_0(\Delta (t),\Delta (s))$  satisfies a differential equation (App. H)
$$
\left(\frac {\partial}{\partial t} + \lambda_ {k}\right) \left(\frac {\partial}{\partial s} + \lambda_ {\ell}\right) \Sigma_ {k \ell} ^ {\Delta} (t, s) = \sum_ {n m} \kappa_ {k m \ell n} \Delta_ {m} ^ {\infty} (t) \Delta_ {n} ^ {\infty} (s), \tag {7}
$$
where  $\Delta_k^\infty (t)\equiv \exp (-\lambda_kt)\langle \psi_k(\pmb {x})y(\pmb {x})\rangle_{\pmb{x}}$  are the errors at infinite width for eigenmode  $k$
(a) Average Prediction Errors