A Dynamical Model of Neural Scaling Laws

Introduction

Neural scaling laws play a central role in modern deep learning practice, and have substantial implications for the optimal trade-off between model size and training time (Hoffmann et al., 2022), as well as architecture selection (Alabdulmohsin et al., 2023). Understanding the origin of such scaling laws, as well as their exponents, has the potential to offer insight into better architectures, the design of better datasets (Sorscher et al., 2022), and the failure modes and limitations of deep learning systems. Yet, many questions about neural scaling laws remain open.

In this paper, we introduce and analyze a solvable model which captures many important aspects of neural scaling laws. In particular, we are interested in understanding the following empirically observed phenomena:

In many domains of deep learning, the test loss of a model with $N$ trainable parameters trained for $t$ iterations has been found to scale as $\mathcal{L}(t,N,P)\approx\mathcal{L}_{0}+a_{t}t^{-r_{t}}+a_{N}N^{-r_{N}}$ (Kaplan et al., 2020; Hoffmann et al., 2022). These scaling law exponents $r_{t},r_{N}$ generally depend on the dataset and architecture. We demonstrate scaling laws on simple vision and language tasks in Figure 1. The compute is proportional to the number of steps of gradient descent times the model size $C\propto Nt$. Setting $N$ and $t$ optimally gives that test loss scales as a power law in $C$. This is the compute optimal scaling law.

Hoffmann et al. (2022) observed that $r_{t}$ and $r_{N}$ are close but slightly different, leading to asymmetric compute-optimal scaling of parameters. For compute budget $C$, they scale model size $N\propto C^{c_{1}}$ and training time $t\propto C^{c_{2}}$ with $c_{2}>c_{1}$. This difference in exponents led to a change in the scaling rule for large language models, generating large performance gains.

Provided feature learning is held constant across model scales (i.e. adopting mean-field or $\mu$P scaling), wider networks tend to train faster (Yang et al., 2021) (Figure 1). If training proceeds in an online/one-pass setting where datapoints are not repeated, then the wider models will also obtain lower test loss at an equal number of iterations. This observation has been found to hold both in overparameterized and underparameterized regimes (Bordelon & Pehlevan, 2023; Vyas et al., 2023).

Early training can be well described by the learning curves for stochastic gradient descent without reuse of samples (termed the online/ideal limiting dynamics), however over time the effect of reusing data accumulates and leads to worse test performance (Nakkiran et al., 2021b; Mignacco et al., 2020; Ghosh et al., 2022). Similarly the gaps in model performance across various model sizes also grow with training time (Yang et al., 2021; Vyas et al., 2023). Figure 1 (d) shows overfitting and reversal of “wider is better” phenomenon due to data reuse.

Prior works (Dyer & Gur-Ari, 2020; Atanasov et al., 2023; Roberts et al., 2022; Bordelon & Pehlevan, 2023) predict early-time finite-width loss corrections that go as $1/\textit{width}$ near the infinite width limit in either lazy or feature-learning regimes. Bahri et al. (2021) et al provide experiments demonstrating the $1/\textit{width}$ convergence. However, finite-width models trained for a long time exhibit non-trivial exponents with respect to model width (Kaplan et al., 2020; Vyas et al., 2023). See Figure 1 for examples of nontrivial scalings at late time on CIFAR-5M and Wikitext.

Near the limit of infinite width, finite models can be thought of as noisy approximations of the infinite-width model with noise that can be eliminated through ensembling (Dyer & Gur-Ari, 2020; Geiger et al., 2020; Atanasov et al., 2023). However recent experiments (Vyas et al., 2023) indicate that ensembling is not enough to match performance of larger models.

These phenomena are not unique to deep networks, but can be observed in linear models, or linearized neural networks operating in the lazy/kernel regime. Though this regime does not capture feature learning, it has benefit of analytical tractability. In this paper, we focus on such linearized models to attempt to gain insight into the dynamics of training.

To attempt to explain these phenomena, we develop a mathematically tractable model of neural scaling laws which allows one to simultaneously vary time, model size, and dataset size. Our contributions are as follows:

We analyze the learning dynamics of a structured and randomly projected linear model trained with gradient descent or momentum. In an asymptotic limit of the model, we obtain a dynamical mean field theory (DMFT) description of the learning curve in terms of correlation functions, which measure the cross-time correlation of training and test errors, and response functions which measure sensitivity of the dynamics to small perturbations.

We solve for the response functions exactly in Fourier domain. This solution reveals faster training for larger models. The low frequency range of these functions allow us to extract the long time limit of the loss.

We show that the model and data corrections to the dynamics accumulate over time. At early time, each of these corrections has a universal scaling, consistent with prior works (Bahri et al., 2021).

For power-law structured features we show that the model exhibits power law scaling of test loss with time, model size and dataset size. While the data and model exponents are the same, the time and model exponents are different in general. We show that this gives rise to an asymmetric compute optimal scaling strategy where training time increases faster than model size.

Our theory explains why ensembling is not compute optimal as it gives less benefit to performance than increase in model size.

We observe in Section 5.1 that feature learning networks can obtain better power law scalings, leading to a better compute optimal frontier. We empirically study this phenomenon in Appendix K.

1 Related Works

The learning curves for linear models with structured (non-isotropic) covariates, including infinite-width kernel regression, have been computed using tools from statistical physics and random matrix theory (Bordelon et al., 2020; Spigler et al., 2020; Canatar et al., 2021; Simon et al., 2021; Bahri et al., 2021; Hastie et al., 2022). Mei & Montanari (2022) analyzed a linear model with random projections of isotropic covariates. There, they study the limiting effects of width and dataset size, and observe model-wise and sample-wise double descent. In (Adlam & Pennington, 2020a) a related model is used to study the finite-width neural tangent kernel (NTK) (Jacot et al., 2018) of a given network. Further, (Adlam & Pennington, 2020b) extends this analysis to understand the different sources of variance in the predictions of random feature models and the effect of ensembling and bagging on the test loss. Other works have extended this to models where an additional untrained projection is applied to the structured covariates (Loureiro et al., 2021, 2022; Zavatone-Veth et al., 2022; Atanasov et al., 2023; Maloney et al., 2022; Zavatone-Veth & Pehlevan, 2023; Ruben & Pehlevan, 2023; Simon et al., 2023). Within this literature, which considered fully trained models, the works of (Bordelon et al., 2020; Spigler et al., 2020) derived power-law decay rates for power-law features which were termed resolution limited by (Bahri et al., 2021) and recovered by (Maloney et al., 2022).

However, we also study the dependence on training time. The $t\to\infty$ limit of our DMFT equations recovers the final losses computed in these prior works. While these prior works find that the scaling exponents for model-size and dataset-size are the same, we find that the test loss scales with a different exponent with training time, leading to a different (model and task dependent) compute optimal scaling strategy.

DMFT methods have been used to analyze the test loss dynamics for general linear and spiked tensor models trained with high-dimensional random data (Mannelli et al., 2019; Mignacco et al., 2020; Mignacco & Urbani, 2022) and deep networks dynamics with random initialization (Bordelon & Pehlevan, 2022b; Bordelon et al., 2023). High dimensional limits of SGD have been analyzed with Volterra integral equations in the offline case (Paquette et al., 2021) or with recursive matrix equations in the online case (Varre et al., 2021; Bordelon & Pehlevan, 2022a). Random matrix approaches have also been used to study test loss dynamics in linear regression with isotropic covariates by (Advani et al., 2020). In this work, we consider averaging over both the disorder in the sampled dataset and the random projection of the features simultaneously.

Other models and hypotheses for scaling laws instead rely on a discrete collection of subtasks or skills which are learned as compute grows (Caballero et al., 2022; Arora & Goyal, 2023; Michaud et al., 2023). Our theory instead focuses on spectral components of a data distribution.

Setup of the Model

We consider a “teacher-student” setting, where data sampled from a generative teacher model is used to train a student random feature model. The teacher and student models mismatch in a particular way that will be described below. This mismatch is the key ingredient that leads to most of the phenomena that we will discuss.

Here $\psi(\bm{x})$ play the role of the infinite-width NTK eigenfunctions, which form a complete basis for square-integrable functions $L^{2}[p]$. The $\epsilon(\bm{x})$ function describes a component of $y$ with which is uncorrelated with $\psi(\bm{x})$. We work in the eigenbasis of features as in (Bordelon et al., 2020), so the covariance given by:

The power law structure in the $\lambda_{k}$ and $\bm{w}^{*}$ entries will lead to power law scalings for the test loss and related quantities.

Our student model is motivated by a scenario where a randomly initialized finite-width network is trained in the linearized or lazy regime (Chizat et al., 2019; Jacot et al., 2018). Such training can be described through learning linear combinations of the finite-width NTK features. These features will span a lower-dimensional subspace of the space of square-integrable functions, and relate to infinite-width NTK features in a complicated way.

Here, we will interpret $N$ as the model size with the $N\to\infty$ limit recovering original kernel. Similar models were studied in (Maloney et al., 2022; Atanasov et al., 2023).

We will focus on the setting where the elements of $\bm{A}$ are drawn iid from a distribution of mean zero and variance one. See Appendix B for details on the technical assumptions. The motivations for this choice are (1) tractability and (2) it satisfies the constraint that as $N\to\infty$ the student’s kernel approaches the infinite-width kernel $\bm{\psi}$. In more realistic settings, such as when projecting the eigenfunctions of an infinite-width NTK to a finite-width NTK, the form of the $\bm{A}$ matrix is generally not known.

The model is trained on a random dataset $\mathcal{D}=\{\bm{x}_{\mu},y_{\mu}\}_{\mu=1}^{P}$ of size $P$ with gradient flow on MSE loss

We explore extensions (momentum, discrete time, one-pass SGD in Appendix J). We track the test and train loss

In small size systems, these losses depend on the precise realization of the data $\mathcal{D}$ and matrix $\bm{A}$. These two quantities can be viewed as the disorder. For large systems, these losses approach a well-defined limit independent of the specific realization of $\mathcal{D},\bm{A}$. We will use this fact in the next section when analyzing the model.

DMFT for Scaling Laws

We next describe a theoretical approach for characterizing the learning curves for this model. The full details of this approach is detailed in Appendices A, B.

We derive a mean field theory for $M,N,P$ large. We analyze both the (1) proportional regime where $N/M=\nu,P/M=\alpha$ and $M,N,P\to\infty$, and (2) non-proportional regime where $M\to\infty$ first and $N,P\gg 1$. The theories derived in these limits are structurally identical (App. F). While the proportional limit is exact, the finite size $N,P$ theory will also contain variability across realizations of disorder. When relevant, we show these in experiments by plotting standard deviations over draws of data and projection matrices $\bm{A}$. This variance decays as $\mathcal{O}(1/P+1/N)$.

The test loss is then given by $\mathcal{L}(t)=\frac{1}{M}\sum_{k}\lambda_{k}v^{0}_{k}(t)^{2}$. The $\bm{v}^{0}$ vector has the following dynamics:

Using DMFT, we characterize this limit by tracking $\bm{v}^{0}$ together with the following random vectors:

The key summary statistics (also called order parameters) are the correlation functions:

Here $\frac{\delta\bm{v}^{i}(t)}{\delta\bm{v}^{j}(s)}$ is the response of $\bm{v}^{i}(t)$ to a kick in the dynamics of $\bm{v}^{j}$ at time $s$. See appendix B.1.1 for details.

The test loss $\mathcal{L}$ and train loss $\hat{\mathcal{L}}$ are related to the time-time diagonal of $C_{0}(t,s)$ and $C_{1}(t,s)$ respectively

These collective quantities concentrate over random draws of the disorder (Sompolinsky & Zippelius, 1981). We show that these correlation and response functions satisfy a closed set of integro-differential equations which depend on $\alpha,\nu$ which we provide in the Appendices A.2.

Further, we show in Appendix A.3 that the response functions possess a time-translation invariance property $R(t,s)=R(t-s)$. This enables exact analysis in the Fourier domain $R(\tau)=\int\frac{d\omega}{2\pi}e^{i\omega\tau}\mathcal{R}(\omega)$. These response functions can then be used to solve for the correlation functions $\{C_{0},C_{1},C_{2},C_{3}\}$.

To understand the convergence of the learned function $f$ along each eigenfunction of the kernel, we introduce the transfer functionThere are dynamical analogues of the mode errors in (Bordelon et al., 2020; Canatar et al., 2021) or learnabilities in (Simon et al., 2021). for mode $k$, $H_{k}(t)\equiv\frac{\left<v_{k}^{0}(t)\right>}{w_{k}^{\star}}$. Our key result is that the Fourier transform of $H_{k}$ can be simply expressed in terms of the Fourier transforms of $R_{1},R_{3}$:

These functions satisfy the self-consistent equations:

From these solved response functions $\mathcal{R}_{1},\mathcal{R}_{3}$, we can compute local solutions to the correlation functions’ two-variable Fourier transform $\mathcal{C}(\omega,\omega^{\prime})$ which are independent equations for each pair of $\omega,\omega^{\prime}$. Information about the early dynamics can be extracted from high frequencies $\omega\gg 1$ while information about the late-time limit of the system can be extracted from $\omega,\omega^{\prime}\to 0$ (App. C, D). For example, for the final test loss,

The full temporal trajectory can be obtained with an inverse Fourier transform of $\mathcal{C}_{0}(\omega,\omega^{\prime})$. See Appendix A.4.

Results

Our results hold for any $\lambda_{k}$ and $w_{k}^{\star}$ and we provide some simple analytically solvable examples in Appendix H. However, based on empirical observations of NTK spectral decompositions on realistic datasets (Bordelon & Pehlevan, 2021; Spigler et al., 2020; Bordelon & Pehlevan, 2022a; Bahri et al., 2021; Maloney et al., 2022), here, we focus on the case of power law features. In this setting, eigenvalues and target coefficients decay as a power law in the index $k$

We will refer to $a$ as the task-power exponent and $b$ as the spectral decay exponentThese power-law decay rates are also known as source and capacity conditions in the kernel literature (Caponnetto & Vito, 2005; Cui et al., 2021). See Figure 6 (a)-(b) for an example with a Residual CNN on CIFAR-5M.

For power law features, the test loss will generally be bottlenecked by either training time $t$ (steps of gradient descent), the size of the training set $P$, or the size of the model $N$. We can derive bottleneck scalings from our exact expressions for $\mathcal{L}(t,P,N)$ (Appendix I):

A consequence of this is an asymmetry in exponent between the model and data bottlenecks compared to the time bottleneck. We verify this asymmetry in Figure 2.

All three of the bottleneck scalings arise due to rank constraints in the effective dynamics. Heuristically, finite training time or the subsampling of data/features leads to an approximate projection of the target function onto the top $k_{\star}(t,P,N)$ eigenspace of the infinite-width kernel. The components of the target function in the null-space of this projection are not learned. This leads to an approximate test loss of the form

For model and data bottlenecks we have that $k_{\star}\propto N$ and $k_{\star}\propto P$ respectively (App. I). On the other hand, $k_{\star}$ for the time bottleneck also depends on the structure of the features through the exponent $b$. This is due to the fact that the $k$-th eigenfeature is learned at a timescale $\tau_{k}\sim k^{b}$. Thus at time $t$, we have learned the first $k_{\star}\approx t^{1/b}$ modes and the variance in the remaining modes gives $\sim t^{-(a-1)/b}$. In the limit of $t\to\infty$ our data and model bottleneck scalings agree with the resolution and variance-limited scalings studied in (Bahri et al., 2021) as well as prior works on kernels and random feature models (Bordelon et al., 2020; Maloney et al., 2022).

These results suggest that in our model, one should scale time and model size differently at compute budget $C=Nt$,

We obtain the above scaling by approximating the loss as a sum of the three terms in equation (14) and a constant, see Appendix M. This analysis suggests that for features that have rapid decay in their eigenspectrum, it is preferable to allocate greater resources toward training time rather than model size as the compute budget increases. This is consistent with the findings of (Hoffmann et al., 2022) for language models. In the limit as $b\to 1$, the time and parameter count should be scaled linearly together. We verify this scaling rule and its $b$-dependence in Figure 3.

Larger models are not always better in terms of test loss for all time $t$, as we showed in Figure 1 (c), especially if the dataset is limited. In Figure 4, we illustrate that larger $N$ can improve convergence to a data-bottlenecked loss for power law features. However, the loss may still be non-monotonic in training time, motivating regularization or early stopping.

The exact gap between train and test losses can exactly be expressed in terms of the DMFT order parameters:

We derive this relation in the Appendix E. At early time this gap goes as $\mathcal{O}(1/P)$ (App. D, E). At late time, however, this picks up a nontrivial task-dependent scaling with $P$ as we show in Figure 2 (e)-(f) and App. C. In Figure 4 (c) we show this gradual accumulation of finite data on the test-train loss gap. For larger datasets $P$ it takes longer training time to begin overfitting (App. E).

Ensembling a set of models means averaging their predictions over the same datasets but with different intitialization seeds. This reduces test loss by reducing the variance of the model output $f$ due to initialization. This improvement can be predicted from an extension of our DMFT (App. G). Analogously, bagging over $B$ datasets reduces variance due to sampling of data.

One might imagine that ensembling many finite sized models would allow one to approach the performance of an infinite sized model ($N\to\infty$). If this were possible, the compute optimal strategy could involve a tradeoff between ensemble count and model size. However, recent experiments show that there is a limited benefit from ensembling on large datasets when compared to increasing model size (Vyas et al., 2023). We illustrate this in Figure 5 (a). Our theory can explain these observations as it predicts the effect of ensembling $E$ times on the learning dynamics as we show in App. G. The main reason to prefer increasing $N$ rather than increasing $E$ is that larger $N$ has lower bias in the dynamics, whereas ensembling only reduces variance. The bias of the model $\mathcal{B}$ has the form

which depend on transfer function $H_{k}$ that we illustrate for power-law features in Figure 5 (b). Since $H_{k}(\tau)$ depend on $N,P$, we see that ensembling/bagging cannot recover the learning curve of the $N,P\to\infty$ system.

Tests on Realistic Networks

We now move beyond synthetic power-law datasets and consider realistic image datasets and architectures. We take the CIFAR-5M dataset introduced in (Nakkiran et al., 2021a) and consider the task of classfiying animate vs inanimate objects. We plot the spectra of the finite-width NTK at initialization across different widths for a Wide ResNet (Zagoruyko & Komodakis, 2016) in Figure 6 a). Here the width parameter corresponds to the number of channels in the hidden layers. Following (Canatar et al., 2021), we define $C(k)$ as the fraction of the task captured by the top $k$ kernel eigenmodes:

Then $1-C(k)$ is the portion of the task left unexplained. We plot this for the initial NTKs across widths in Figure 6 b). We extract the spectral decay exponent $b$ and the the task power exponent $a$ from these two curves. Together, these give the learning scaling laws of the linearized neural network model on this dataset. We plot the compute optimal scaling laws of these linearized models in Figure 6 c). We also plot the predicted scaling law $C^{-(a-1)/(1+b)}$ in blue and find excellent agreement.

We also compare these scalings to those of the compute optimal learning curves for feature-learning networks. We train several networks with different widths and initialization seeds for 64 epochs through the dataset. We observe substantially different compute-optimal scaling exponents in the dotted curves of Figure 6 c). This means that although our random feature model does capture the correct linearized scaling trends, which have all of the qualities observed in realistic scaling laws, more is needed to capture the acceleration of scaling induced by feature learning. Further analyses of the after-kernels of feature learning networks are performed in Appendix K. We see that the kernels continue to evolve substantially throughout training. This indicates that a full explanation of the compute optimal scaling exponents will require something resembling a mechanistic theory of kernel evolution (Long, 2021; Fort et al., 2020; Atanasov et al., 2022; Bordelon & Pehlevan, 2022b).

Conclusion

We have presented a model that recovers a wide variety of phenomena observed in more realistic deep learning settings. Our theory includes not just model size and dataset size as parameters but also explicitly treats the temporal dynamics of training. We observe different scaling exponents for performance in terms of model size and number of time steps. Future work to incorporate kernel evolution into this model could further shed insight into the improved scaling laws in the feature-learning regime. Overall, our results provide a theoretical interpretation of compute-optimal scaling as a competition between the training dynamics of the infinite width/infinite data limit and finite model-size bottleneck.

Broader Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

Acknowledgements

We are grateful to Yasaman Bahri, Stefano Mannelli, Francesca Mignacco, Jascha Sohl-Dickstein, and Nikhil Vyas for useful conversations. We thank Clarissa Lauditi and Jacob Zavatone-Veth for comments on the manuscript.

B.B. is supported by a Google PhD Fellowship. A.A. is supported by a Fannie and John Hertz Fellowship. C.P. is supported by NSF grant DMS-2134157, NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.

References

Appendix A Derivation of Dynamical Model of Scaling Laws

We investigate the simplest possible model which can exhibit task-dependent time, model size and finite data bottlenecks. We therefore choose to study a linear model with projected features

The weights $\bm{w}$ are updated with gradient descent on a random training dataset which has (possibly) noise corrupted target values $y_{\mu}=y(\bm{x}_{\mu})+\sigma\epsilon_{\mu}$. This leads to the following gradient flow dynamics

We introduce the variable $\bm{v}^{0}=\bm{w}_{\star}-\frac{1}{\sqrt{N}}\bm{A}^{\top}\bm{w}$ to represent the residual error of the learned weight vector. This residual error has the following dynamics.

The entries of each matrix are treated as random with $\Psi^{\mu}_{k}\sim\mathcal{N}(0,\lambda_{k})$ and $A_{jk}\sim\mathcal{N}(0,1)$. To study the dynamical evolution of the test error $\mathcal{L}(t)=\frac{1}{M}\bm{v}^{0}(t)^{\top}\bm{\Lambda}\bm{v}^{0}(t)+\sigma^{2}$, we introduce the sequence of vectors

The train and test losses can be computed from the $\bm{v}^{0}$ and $\bm{v}^{1}$ fields

In the next section, we derive a statistical description of the dynamics in an appropriate asymptotic limit using dynamical mean field theory methods.

Standard field theoretic arguments such as the cavity or path integral methods can be used to compute the effective statistical description of the dynamics in the limit of large $M,N,P$ with fixed ratios $\alpha=P/M$ and $\nu=\frac{N}{M}$ (see Appendix B). This computation gives us the following statistical description of the dynamics.

The correlation and response functions obey

These equations are exact in the joint proportional limit for any value of $\alpha,\nu$.

A.2 Closing the Equations for the Order Parameters

Though we expressed the dynamics in terms of random fields, we stress in this section that all of the dynamics for the correlation and response functions close in terms of integro-dfferential equations. To shorten the expression, we will provide the expression for $\beta=0$, but momementum can easily be added back by making the substitution $\partial_{t}\to\partial_{t}+\beta\partial^{2}_{t}$.

First, our closed integral equations for the response functions are

We note that these equations imply causality in all of the response functions since $R(t,s)=0$ for $t<s$. Once these equations are solved for the response functions, we can determine the correlation functions, which satisfy

Solving these closed equations provide the complete statistical characterization of the limit. The test and train losses are given by the time-time diagonal of $C_{0}(t,t),C_{1}(t,t)$.

A.3 Time-translation Invariant (TTI) Solution to Response Functions

From the structure of the above equations, the response functions are time-translation invariant (TTI) since they are only functionals of TTI $\delta(t-s)$ Dirac-Delta function and $\Theta(t-s)$ Heaviside step-function. As a consequence, we write each of our response functions in terms of their Fourier transforms

We will keep track of the regulator $\epsilon$ and consider $\epsilon\to 0^{+}$ at the end of the computation. The resulting DMFT equations for the response functions have the following form in Fourier space

where $\epsilon\to 0$ will be taken after. Combining these equations, we arrive at the simple set of coupled equations

After solving these equations for all $\omega$, we can invert the dynamics of $v^{0}_{k}(t)$ to obtain its Fourier transform

where we defined the transfer functions $\mathcal{H}_{k}(\omega)$. From this equation, we can compute $v^{0}_{k}(t)$ through inverse Fourier-transformation and then compute the correlation function to calculate the test error. An interesting observation is that the response functions $\mathcal{R}_{1}(\omega),\mathcal{R}_{3}(\omega)$ alter the pole structure in the transfer function, generating $\nu,\alpha$ dependent timescales of convergence.

A.4 Fourier Representations for Correlation Functions

While the response functions are TTI, the correlation functions transparently are not (if the time-time diagonal $C_{0}(t,t)$ did not evolve, then the loss $\mathcal{L}(t)$ wouldn’t change!). We therefore define the need to define the double Fourier transform $\mathcal{C}(\omega,\omega^{\prime})$ for each correlation function $C(t,s)$

Assuming that all response functions and transfer functions $\mathcal{H}_{k}$ have been solved for, the correlation functions satisfy the closed set of linear equations.

These equations can be efficiently solved for all pairs of $\omega,\omega^{\prime}$ after the response functions have been identified. Then one can take an inverse Fourier transform in both indices.

Appendix B Field Theoretic Derivation of DMFT Equations

In this section, we derive the field theoretic description of our model. We will derive this using both the Martin-Siggia-Rose (MSR) path integral method (Martin et al., 1973) and the dynamical cavity method. For a recent review of these topics in the context of neural networks, see (Helias & Dahmen, 2020).

With the MSR formalism, we evaluate the moment generating functional for the field $\bm{v}^{0}(t)$:

Note that at zero source, we have the important identity that

We insert a Dirac delta functions to enforce the definitions of each of the fields $\{\bm{v}^{1},\bm{v}^{2},\bm{v}^{3},\bm{v}^{4}\}$ as in equation 8.

We now average over the sources of disorder. We assume that the entries of $\bm{A}$ are i.i.d. with mean zero and variance 1. In the proportional limit, we can replace the entries of $\bm{A}$ as a draw from a Gaussian $\mathcal{N}(0,1)$ by appealing to Gaussian equivalence. We furhther justify this in the cavity derivation in the next section. This allows us to evaluate the averages over the matrix $\bm{A}$.

Similarly, we can calculate the averages over the data, which enters via the design matrices $\bm{\Psi}$. Again in this proportional limit we can use a Gaussian equivalence on $\bm{\Psi}$ to have it take the form $\bm{\Psi}\sim\bm{\Phi}\bm{\Lambda}^{1/2}$ where $\bm{\Phi}$ has entries drawn from a unit normal. Taking the average then gives us

We now insert delta functions for following bracketed terms: $C_{0},C_{1},C_{2},C_{3}$ and $R_{1},R_{3}$ using the following identity (e.g. for $C_{0}$ at times $s,t$):

Here the $\hat{C}_{i},\hat{R}_{i}$ integrals are taken over the imaginary axis. This yields a moment generating function (here we’ll take $\bm{j}=0$):

The constraint that $Z=1$ means that $S=0$ at the saddle point. $S$ here is given by:

We have chosen to take $\hat{R}_{i}(s,t)$ to have a different sign and $s,t$ ordering convention than the $\hat{C}_{i}$ to simplify our notation later on. We have also used that Equations (38), (39) factorize over their respective indices, so each $\mathcal{Z}$ is a partition function over a single index. The individual $\mathcal{Z}_{i}$ are given by:

In the large $M$ limit we evaluate this integral via saddle point. The saddle point equations give:

Here $\langle\cdot\rangle$ denotes an average taken with respect to the statistical ensemble given by the corresponding partition function $\mathcal{Z}_{i}$. Following the discussion below Equation 37, we take $\hat{C}_{i}(t,s)=0$, which will enforce $\langle\hat{v}_{i}(t)\hat{v}_{i}(s)\rangle=0$ and lead to the correct dynamical equations.

To evaluate the remaining, we can integrate out the $\hat{v}^{i}$ variables. First let us look at $\mathcal{Z}_{1}$. Using the Hubbard-Stratonovich trick we can write the action in terms linear in $\hat{v}^{1}$. This gives

We now replace $\hat{C}_{1}$ by its saddle point value of and $\hat{R}_{1}$ by $R_{0,2}$. Integrating over $\hat{v}$ gives a delta function:

For $\mathcal{Z}_{0,2,4;k}$ after replacing $\hat{C}_{0},\hat{C}_{2},\hat{C}_{4}$ with their saddle point values we get:

Using the same Hubbard-Stratonovich trick on $\hat{v}^{2}_{k}$ gives:

Lastly, the equations of motion for $v^{0}_{k}$ in terms of $v^{4}_{k}$ are known:

One can easily add momentum by replacing $\partial_{t}v^{0}_{k}(t)$ with $(\beta\partial^{2}_{t}+\partial_{t})v^{0}_{k}(t)$ without changing anything else about the derivation.

B.2 Cavity Derivation

The cavity derivation relies on Taylor expanding the dynamics upon the addition of a new sample or feature. We will work through each cavity step one at a time by considering the influence of a single new base feature, new sample, and new projected feature. In each step, the goal is to compute the marginal statistics of the added variables. This requires tracking the linear response to all other variables in the system.

The next order corrections have a subleading influence on the dynamics. Now, inserting these perturbed dynamics into the dynamics for the new $(M+1)$st set of variables $\{v^{2}_{0}(t),v^{4}_{0}(t)\}$. For $v^{2}_{0}(t)$, we have

There are now two key steps in simplifying the above expression in the proportional limit:

By the fact that the $v^{1}_{\mu}(t)$ dynamics are statistically independent of the new feature $\psi^{\mu}_{0}$, we can invoke a central limit theorem for the first term which is mean zero and variance $\mathcal{O}(1)$.

Similarly, we can invoke a law of large numbers for the second term, which has $\mathcal{O}(1)$ mean and variance on the order of $\mathcal{O}(M^{-1})$. Therefore in the asymptotic limit it can be safely approximated by its mean.

We note in passing that neither of these steps require the $\psi^{\mu}_{0}$ variables to be Gaussian. Thus we obtain the following asymptotic statistical description of the $v^{2}_{0}(t)$ random variable

Following an identical argument for $v^{4}_{0}(t)$ we have

Now, computing the dynamics of the new random variable $v^{1}_{0}(t)$

Now, we finally consider the effect of introducing a single new projected feature so that instead of $N$ we now have $N+1$ projected features. This causes a perturbation to $\{v^{2}_{k}(t)\}$ which we

Now, we compute the dynamics for the added variable $v^{3}_{0}(t)$

Now, using the information gained in the previous sections, we can combine all of the dynamics for each field into a closed set of stochastic processes. This recovers the DMFT equations of Appendix A.2.

Appendix C Final Losses (the t→∞→𝑡t\to\infty Limit of DMFT)

In this section we work out exact expressions for the large time limit of DMFT. By comparing with prior computations of the mean-field statics of this problem computed in (Atanasov et al., 2023; Zavatone-Veth & Pehlevan, 2023; Ruben & Pehlevan, 2023; Maloney et al., 2022; Simon et al., 2021), we show that the large time and large $M$ limits commute, specifically that $\lim_{M,N,P\to\infty}\lim_{t\to\infty}\mathcal{L}(M,N,P,t)=\lim_{t\to\infty}\lim_{M,N,P\to\infty}\mathcal{L}(M,N,P,t)$. We invoke the final value theorem and use the response functions as before.

We note that for functions which vanish at $t=-\infty$, that

where we invoked integration by parts and used the assumption that $\lim_{\tau\to-\infty}H(\tau)=0$, a condition that is satisfied for the correlation and response functions in our theory. We can therefore use the identity $\lim_{\tau\to\infty}H(\tau)=\lim_{\omega\to 0}i\omega\mathcal{H}(\omega)$ to extract the final values of our order parameters.

We also need to invoke a similar relationship for the final values of the correlation functions

where $\mathcal{C}$ is the two-variable Fourier transform. The final value of the test loss is $\lim_{t\to\infty}\mathcal{L}(t)=\lim_{t,s\to\infty}C_{0}(t,s)$.

C.1 General Case (Finite ν,α𝜈𝛼\nu,\alpha)

Before working out the solution to the response functions, we note that the following condition is always satisfied

For $\nu=\alpha$, this equation implies that $\mathcal{R}_{1}=\mathcal{R}_{3}$. For $\nu\neq\alpha$, we can have either $\mathcal{R}_{1}\to 0$ or $\mathcal{R}_{3}\to 0$ but not both. We consider each of these cases below.

In this case, the response function $\mathcal{R}_{1}\sim\mathcal{O}(i\omega)$ as $\omega\to 0$ and $\mathcal{R}_{3}\sim 1-\frac{\alpha}{\nu}$ as $\omega\to 0$. We thus define

Using the equation which defines $\mathcal{R}_{1}$, we find that the variable $r$ satisfies the following relationship at $\omega\to 0$

After solving this implicit equation, we can find the limiting value of $i\omega\mathcal{H}_{k}(\omega)$ as

Next, we can work out the scaling of the correlation functions in the limit of low frequency. We define the following limiting quantities based on a scaling analysis performed on our correlation functions for small $\omega$

These limiting quantities satisfy the closed set of linear equations

These equations can be solved for $\{C_{0}^{\infty},C_{1}^{\infty},C_{2}^{\infty},C_{3}^{\infty}\}$. Simplifying the expressions to a two-variable system, we find

This expression recovers the ridgeless limit of the replica results of (Atanasov et al., 2023; Zavatone-Veth & Pehlevan, 2023) and the random matrix analysis of (Simon et al., 2023).

Following the same procedure, we note that for $\nu<\alpha$ that $\mathcal{R}_{3}\sim\mathcal{O}(i\omega)$ and $\mathcal{R}_{1}\sim 1-\frac{\nu}{\alpha}$. We thus find the following equation for $r=\lim_{\omega\to\infty}(i\omega)^{-1}\mathcal{R}_{1}(\omega)\mathcal{R}_{3}(\omega)$.

where as before $H_{k}^{\infty}=\frac{1}{1+\lambda_{k}r}$. The analogous scaling argument for small $\omega$ gives us the following set of well-defined limiting quantities

where these limiting correlation values satisfy

This is again a closed linear system of equations for the variables $\{C_{0}^{\infty},C_{1}^{\infty},C_{2}^{\infty},C_{3}^{\infty}\}$. In the next section, we recover the result for kernel regression where $\nu\to\infty$ and the learning curve for infinite data $\alpha\to\infty$ with respect to model size $\nu$.

C.2 Learning Curves for Kernel Regression ν,t→∞→𝜈𝑡\nu,t\to\infty

In the $t\to\infty$ and $\nu\to\infty$ limit we recover the learning curve for kernel regression with eigenvalues $\lambda_{k}$. To match the notation of (Canatar et al., 2021), we define

which generates the following self-consistent equation for $\kappa$

Plugging this into the expression for the loss, we find

Letting $C_{\infty}\equiv\lim_{\omega,\omega^{\prime}\to 0}(i\omega)(i\omega^{\prime})\mathcal{C}(s,s^{\prime})$, we have

The variable $\kappa$ decreases from $\left[\frac{1}{M}\sum_{k}\lambda_{k},0\right]$ as $\alpha\in$. For $\alpha>1$ we have $\kappa=0$. The quantity $\frac{1}{1-\gamma}$ comes from overfitting due to variance from the randomly sampled dataset.

Appendix D Early Time Dynamics (High-Frequency Range)

In this section, we explore the early time dynamical effects of this model. Similar to how the late time dynamical effects could be measured by examining the low frequency $\omega\ll 1$ part of the response and correlation functions, in this section, we analyze the high frequency components $\omega\gg 1$. We start by noting the following expansions valid near $\omega\to\infty$

We let $c=\frac{1}{M}\sum_{k}\lambda_{k}$. These can be plugged into the transfer function for mode $k$

Preforming an inverse Fourier transform, we find the following early time asymptotics

We see from this expression that the early time corrections always scale as $1/\alpha$ or $1/\nu$ and that these corrections build up over time. A similar expansion can be performed for $\mathcal{C}(\omega,\omega^{\prime})$ with $\omega,\omega^{\prime}\gg 1$ which also gives leading corrections which scale as $1/\alpha$ and $1/\nu$.

Appendix E Buildup of Overfitting Effects

In this section, we derive a formula for the gap between test loss $\mathcal{L}(t)$ and train loss $\hat{\mathcal{L}}(t)$. We start from the following formula

Moving the $v_{1}(t)$ term to the other side, and using the fact that $\left<u_{1}(t)u_{1}(s)\right>=C_{1}(t,s)$, we find the following relationship between train and test loss

To get a sense of these expressions at early and late timescales, we investigate the Fourier transforms at high $\omega\gg 1$ and low $\omega\ll 1$ frequencies respectively.

The relationship between Fourier transforms at high frequencies $\omega\gg 1$ is

where $c=\frac{1}{M}\sum_{k}\lambda_{k}$. Taking a Fourier transform back to real time gives us the following early time differential equation for the test-loss train loss gap

The above equation should hold for early times. We note that $C_{0}(t,t)-C_{1}(t,t)=\mathcal{L}(t)-\mathcal{\hat{L}}(t)$ exactly recovers the test-train gap.

E.2 Low Frequency Range/Late Time

At late time/low frequency, as we showed in Appendix C, the behavior of the $C_{1}$ correlation function depends on whether the model is over-parameterized or under-parameterized. In the overparameterized case, the asymptotic train loss is zero while the asymptotic test loss is nonzero. In the underparameterized case, we have a limiting value for both the test and train loss which can be computed from the expressions in Appendix C.

Appendix F Non-Proportional (Dimension-Free) Limit

We can imagine a situation where the original features are already infinite dimensional ($M\to\infty$ is taken first). This would correspond more naturally to the connection between infinite dimensional RKHS’s induced by neural networks at infinite width (Bordelon et al., 2020; Canatar et al., 2021; Cheng & Montanari, 2022). Further, we will assume a trace class kernel $K(\bm{x},\bm{x}^{\prime})=\bm{\psi}(\bm{x})\cdot\bm{\psi}(\bm{x}^{\prime})$ for the base features $\bm{\psi}$ which diagonalizes over the data distribution $p(\bm{x})$ as

As before, we are concerned with the test and train losses

The appropriate scaling of our four fields of interest in this setting are

Following the cavity argument given in the previous section, we can approximate the the correlation and response functions as concentrating to arrive at the following field description of the training dynamics

which are exactly the same equations as in the proportional limit except with the substitution $\nu\to N$ and $\alpha\to P$. The correlation and response functions have the form

which will all be $\mathcal{O}(1)$ under this scaling.

Appendix G Effect of Ensembling and Bagging on Dynamics

We adopt the language of the fine-grained bias-variance decomposition in (Adlam & Pennington, 2020b). There, a given learned function generally depends on both the dataset $\mathcal{D}$ and initialization seed $\theta_{0}$. We write this as $f_{\mathcal{D},\mathcal{\theta}_{0}}$. The role of random initialization is played by the $\bm{A}$ matrix in our setting. For a given function, its variance over datasets and its variance over initializations are respectively given by

G.2 Derivation

In this section, we consider the effect of ensembling over $E$ random initial conditions and bagging over $B$ random datasets. We let $\bm{v}^{0}_{e,b}(t)$ represent the weight discrepancy for model $e$ on dataset $b$. Here $e$ runs from $1$ to $E$ and $b$ runs from $1$ to $B$. The $(e,b)$th vector has dynamics:

Ensembling and bagging would correspond to averaging these $\bm{v}^{0}$s over these $EB$ systems

The key vectors to track for this computation are

We can further show that the $\bm{v}^{0}_{e,b}$ and $\bm{v}^{0}_{e^{\prime},b^{\prime}}$ have response functions that decouple across $e,b$. Intuitively, giving the dynamical system $e,b$ a kick should not alter the trajectory of the separate $e^{\prime},b^{\prime}$ dynamical system, even if they share disorder $\{\bm{\Psi},\bm{A}\}$. The DMFT description of the proportional limit yields the following integral equations for the $v$ fields:

Here, the response functions $R$ are to be computed within a single system. In what follows, we will use $\langle\cdot\rangle$ to denote averages over the disorder, and explicitly write out any averages over the ensemble members and datasets.

The Gaussian variables in the DMFT have the following covariance

The covariances above $C_{0,e,e^{\prime}},C_{1,e,e^{\prime}},C_{2,b,b^{\prime}},C_{3,b,b^{\prime}}$ allow for different ensemble or dataset index but not both. We will use $C_{0},C_{1},C_{2},C_{3}$ etc to represent the correlation functions within a single system. For instance, $C^{0}_{e,e^{\prime}}(t,s)=\frac{1}{M}\sum_{k}\lambda_{k}\left<v^{0}_{e,b}(t)v^{0}_{e^{\prime},b}(s)\right>$ while $C^{0}=\frac{1}{M}\sum_{k}\lambda_{k}\left<v^{0}_{e,b}(t)v^{0}_{e,b}(s)\right>$. The correlation function of interest is thus

We can combine the first two equations and the second two equations to identify the structure of the cross-ensemble and cross-dataset (across-system) correlations in terms of the marginal (within-system) correlation statistics

These equations give the necessary cross-ensemble and cross-dataset correlations. Now we can consider the effect of ensembling and bagging on the dynamics. To do so, consider the Fourier transform of the bagged-ensembled error $\bar{v}^{0}_{k}(t)=\frac{1}{EB}\sum_{eb}v^{0}_{k,e,b}(t)$, which has the Fourier transform

Computing the correlation function for this bagged-ensembled field random variable, we find

The first term is the irreducible bias for mode $k$ which is the loss for mode $k$ when the learned function is averaged over all possible datasets and all possible projections. We see that the second term scales as $\frac{1}{\nu E}$ which will persist even if $B\alpha\to\infty$. Similarly, there is a term that is order $\frac{1}{\alpha B}$ which will persist even if $\nu E\to\infty$. Lastly, there are two terms which depend on both $B,E$. This is similar to the variance that is explained by the interaction of the dataset and the random projection (Adlam & Pennington, 2020b). The test loss is then a Fourier transform of the above function

If $E,B\to\infty$, then we obtain the stated irreducible bias of the main paper

This is the error of the mean output function over all possible datasets and random projections of a certain size.

G.3 Ensembling is Not Always Compute Optimal

For a compute budget $C=NEt$, we find that ensembling does not provide as much benefit as increasing the size of the model. From the results in the last section, we note that ensembling reduces the variance. For this section, we consider the $P\to\infty$ limit. We let $\mathcal{B}(N,t)$ represent the bias and $\mathcal{V}(N,t)$ represent the variance within a single ensemble. The loss at fixed compute then takes the form

For any $\nu$ which satisfies the condition that

we have that ensembling is strictly dominated by increasing $\nu$.

Appendix H White Bandlimited Model

To gain intuition for the model, we can first analyze the case where $\lambda_{k}=1$, which has a simpler DMFT description since each of the $M$ features are statistically identical. We illustrate the dependence of the loss on model size $\nu$ and training time $t$ for $\alpha<1$ in Figure 7. We note that the loss can be non-monotonic in $\nu$ at late training times, but that monotonicity is maintained for optimal early stopping, similar to results on optimal regularization in linear models (Advani et al., 2020) and random feature models (Mei & Montanari, 2022; Simon et al., 2023).

In the case of all $\lambda_{k}=1$ we have the following definitions

Writing $\mathcal{R}_{1}=1-\frac{\nu}{\alpha}(\mathcal{R}_{3}-1)$ allows us to solve for $\mathcal{R}_{3}$ exactly:

This is a cubic equation that can be solved for $\mathcal{R}_{3}$ as a function of $\omega$. In the limit of $\alpha\to\infty$ this simplifies to:

H.2 Timescale Corrections in The Small ν𝜈\nu Regime

By expanding the above in the limit of small $\nu$ we get that $\mathcal{R}_{3}$ goes as

From this approximate response function, we find that the transfer function takes the form

where in the last line, we used the residue theorem. We note that in this perturbative approximation that this transfer function is always greater than the transfer function at $\nu\to\infty$ which is $e^{-\tau}$. Thus finite $\nu$ leads to higher bias in this regime. We define bias and variance precisely in Appendix G.1.

H.3 Timescale corrections in fully expressive regime ν>1𝜈1\nu>1

For $\nu\gg 1$, we can approximate $R_{3}(\omega)\sim 1-\nu^{-1}(1+i\omega)^{-1}$, we have

where we used the residue theorem after closing the contour in the upper half-plane. In Figure 8, we show that this perturbative approximation does capture a slowdown in the dynamics for large but finite $\nu$.

Appendix I Power-Law Bottleneck Scalings

In this section we calculate the scaling of the loss with the various limiting resources (time, model size, and data) when using power law features. Since the power-law features give a trace class kernel (i.e. $\sum_{k=1}^{\infty}\lambda_{k}<\infty$), we use the non-proportional limit formalism in Appendix F, which gives an expression for $\mathcal{L}(t,N,P)$ with $M$ already considered infinite. While the resulting expressions are not a formal proportional thermodynamic limit and finite $N,P$ corrections exist in the form of fluctuations from one random realization of the system to another. These corrections decay rapidly enough at finite $N,P$ for this mean field theory to be accurate and descriptive in realistic systems (Bordelon et al., 2020; Simon et al., 2023; Cheng & Montanari, 2022). We plot this variability of random finite size experiments as highlighted standard deviations in the main text figures.

The time bottleneck is defined as the limiting dynamics in the absence of any model or data finite size effects. To eliminate those effects, we simply study the $\alpha,\nu\to\infty$ limit

In this limit, the response functions simplify to $\mathcal{R}_{1}(\omega)\mathcal{R}_{3}(\omega)\to 1$ so that

Further, in this limit, we have that $C_{0}(t,s)=\frac{1}{M}\sum_{k}\lambda_{k}H_{k}(t)H_{k}(s)(w^{\star}_{k})^{2}$ since all the variance terms (which depend on $\nu^{-1},\alpha^{-1}$) drop out. Thus we have the following loss at time $t$,

where the final scaling with time can be obtained through either change of variables or steepest descent methods (Bordelon & Pehlevan, 2022a).

I.2 Model Bottleneck

In this section we take $\alpha,t\to\infty$. This leaves us with the following equation for $r\equiv\lim_{\omega\to 0}(i\omega)^{-1}\mathcal{R}_{3}(\omega)$.

Now, the large time limit of the transfer functions $H_{k}(\tau)$ can be obtained from the final-value theorem

Now, integrating over the eigenvalue density to get the total loss gives

Thus we expect a powerlaw scaling of the form $\mathcal{L}\sim N^{-(a-1)}$ in this regime.

I.3 Data Bottleneck

In this section we take $\nu,t\to\infty$. This leaves us with the following equation for $r\equiv\lim_{\omega\to 0}(i\omega)^{-1}\mathcal{R}_{1}(\omega)$.

Now, the large time limit of the transfer functions $H_{k}(\tau)$ can again be obtained from the final-value theorem

Now, integrating over the eigenvalue density to get the total loss gives

The loss will therefore scale as $P^{-(a-1)}$ in this data-bottleneck regime.

Appendix J Optimization Extensions

In this section, we point out that DMFT can also completely describe discrete time training as well. In this section we consider discrete time gradient descent with learning rate $\eta$

The delta function in this context is defined as

ensures that the initial condition $v_{k}^{0}(0)=w^{\star}_{k}$ is satisfied. These iteration equations can be closed for the response functions and correlation functions and solved over $T\times T$ matrices.

Alternatively, we can also solve this problem in an analogous frequency space. Analogous to the Fourier transform method, the equations in discrete time can be closed in terms of the $Z$-transform

Applying this transform gives us the following expression for the $v^{0}_{k}$ fields.

Similar to the Fourier case, the final losses can be extracted as the $z\to 1$ limit of these objects.

J.2 Momentum

As mentioned in appendix B, it is straightforward to extend the DMFT treatment beyond just gradient descent dynamics to include a momentum term with momentum $\beta$.

We first consider this replacement in continuous time. This requires applying the following replacement:

This slightly modifies the expressions for the response functions. For example, in Fourier space the response functions become:

In discrete time, momentum updates can be expressed as

where $\bm{b}(t)$ is the filtered version of the loss gradient (the $\bm{v}^{4}(t)$ field) with momentum coefficient $\mu$ and $\eta$ is the learning rate. The dependence on the $\bm{b}(t)$ field can be eliminated by turning this into a second order difference equation

Again, the final result can be expressed in terms of the $Z$-transformed transfer functions $\mathcal{H}_{k}(z)$ which have the form

J.3 One Pass SGD

In this section we derive online SGD with projected features. At each step a random batch of $P=\alpha M$ samples are collected (independent of previous samples), giving a matrix $\bm{\Psi}(t)$ of sampled features. The update at step $t$ is

The DMFT limit gives the following statistical description of the fields, which decouple over time for the $v^{1}(t),v^{2}_{k}(t)$ but remain coupled across time for $v^{3}(t),v^{4}_{k}(t)$

This system cannot exhibit overfitting effects as we have the statistical equivalence between the covariance of $v^{1}$ and the test loss:

We note that this is very different than the case where data is reused at every step, which led to a growing gap between train and test loss as we derive in Appendix E.

Appendix K Kernel Analysis of Feature Learning Networks

In Section 5.1, we observed that feature learning networks can achieve better loss and compute-optimal scaling. In such settings, it may be useful to observe the after kernel, namely the NTK at the end of training. This object can often shed insight into the structure of the learned network function (Atanasov et al., 2022; Long, 2021) and its generalization. In some cases, it has been observed that the final kernel stabilizes during the course of training (Fort et al., 2020), potentially allowing one to potentially deduce scaling laws from the spectrum and task-model alignment of this after-kernel, though other papers have observed contrary results (Vyas et al., 2022).

Motivated by this, we study the NTKs of the finite-width networks trained for 64 epochs with the animate-inanimate CIFAR-5m discrimination task. We observe in Figure 9 a) that the spectrum becomes flatter, with a decay exponent of close to $1.4$ down from $2.0$ for the initial kernel.

The fraction of the task power unexplained is also observed to have a lower exponent in Figure 9 b), however there is also the presence of a low rank spike indicative of the kernel aligning to this discrimination tasks.

From these scalings we can obtain the $a$ and $b$ exponents and get a prediction for the scaling of the test loss. We plot this in grey in Figure 9 c). The observed scaling (in black) is much better than that predicted by the after-kernel. This is an indication the the after kernel continues evolving in this task, improving the scaling exponent of the test loss.

The kernel-target alignment (Cortes et al., 2012), as measured by

is plotted in 9 d). Here $\bm{y}$ is the target labels on a held-out test set, and $\bm{K}$ is the gram matrix of the after-kernel on this test set. We indeed observe a consistent increase in this quantity across time. This gives an indication that understanding the evolution of the after-kernel will be useful

Appendix L Numerical Recipes

The simplest way to solve the DMFT equations is to iterate them from a reasonable initial condition (Mignacco et al., 2020; Bordelon & Pehlevan, 2022b). We solve in discrete time for $T\times T$ matrices $\{\bm{R}_{0,2},\bm{R}_{1},\bm{R}_{2,4},\bm{R}_{3},\bm{C}_{0},\bm{C}_{1},\bm{C}_{2},\bm{C}_{3}\}$ which have entries $[\bm{R}]_{t,s}=R(t,s),[\bm{C}]_{t,s}=C(t,s)$, etc. We let $\bm{\Theta}(t,s)=\eta\Theta(t-s)$ where $\eta$ is the learning rate.

Solve for the response functions by updating the closed equations as matrices by iterating the equations.

Once these response functions have converged, we can iterate the equations for the correlation functions

After iterating these equations, one has the discrete time solution to the DMFT order parameters and any other observable can then be calculated.

L.2 Fourier Transform Method

To accurately compute the Fourier transforms in the model/data bottleneck regime ($\alpha<1$ or $\nu<1$) we have that $\mathcal{R}_{1}(\omega)\mathcal{R}_{3}(\omega)\sim i\omega r$ as $\omega\to 0$ so we must resort to analyzing the principal part and the delta-function contribution to the integral. Construct a shifted and non-divergent version of the function $\mathcal{H}(\omega)$.

The Dirac mass is trivial to integrate over giving $\frac{1}{2(1+r)}$. Lastly, we must perform an integral of the type

Adding these two terms together, our transfer function has the form

The last integral can be performed numerically, giving a more stable result.

Appendix M Compute Optimal Scaling from Sum of Power-Laws

We suppose that the loss scales as (neglecting irrelevant prefactors)

Our goal is to minimize the above expression subject to the constraint that compute $C=Nt$ is fixed. Since $C$ is fixed we can reduce this to a one-dimensional optimization problem

The optimality condition $\partial_{N}L=0$ is

From this last expression one can evaluate the loss at the optimum