How Feature Learning Can Improve Neural Scaling Laws

## Abstract

Abstract We develop a solvable model of neural scaling laws beyond the kernel limit. Theoretical analysis of this model shows how performance scales with model size, training time, and the total amount of available data. We identify three scaling regimes corresponding to varying task difficulties: hard, easy, and super easy tasks. For easy and super-easy target functions, which lie in the reproducing kernel Hilbert space (RKHS) defined by the initial infinite-width Neural Tangent Kernel (NTK), the scaling exponents remain unchanged between feature learning and kernel regime models. For hard tasks, defined as those outside the RKHS of the initial NTK, we demonstrate both analytically and empirically that feature learning can improve scaling with training time and compute, nearly doubling the exponent for hard tasks. This leads to a different compute optimal strategy to scale parameters and training time in the feature learning regime. We support our finding that feature learning improves the scaling law for hard tasks but not for easy and super-easy tasks with experiments of nonlinear MLPs fitting functions with power-law Fourier spectra on the circle and CNNs learning vision tasks.

## 1 Introduction

Deep learning models tend to improve in performance with model size, training time and total available data. The dependence of performance on the available statistical and computational resources are often regular and well-captured by a power-law (Hestness et al., ; Kaplan et al., ). For example, the Chinchilla scaling law (Hoffmann et al., ) for the loss $\mathcal{L}(t,N)$ of a $N$-parameter model trained online for $t$ steps (or $t$ tokens) follows

$$ $\displaystyle\mathcal{L}(t,N)=c_{t}t^{-r_{t}}+c_{N}N^{-r_{N}}+\mathcal{L}_{ \infty},$ (1) $$

where the constants $c_{t},c_{N}$ and exponents $r_{t},r_{N}$ are dataset and architecture dependent and $\mathcal{L}_{\infty}$ represents the lowest achievable loss for this architecture and dataset. These scaling laws enable intelligent strategies to achieve performance under limited compute budgets (Hoffmann et al., ) or limited data budgets (Muennighoff et al., ). A better understanding of what properties of neural network architectures, parameterizations and data distributions give rise to these neural scaling laws could be useful to select better initialization schemes, parameterizations, and optimizers (Yang et al., ; Achiam et al., ; Everett et al., ) and develop better curricula and sampling strategies (Sorscher et al., ).

Despite significant empirical research, a predictive theory of scaling laws for deep neural network models is currently lacking. Several works have recovered data-dependent scaling laws from the analysis of linear models (Spigler et al., ; Bordelon et al., ; Bahri et al., ; Maloney et al., ; Simon et al., ; Bordelon et al., ; Zavatone-Veth & Pehlevan, ; Paquette et al., ; Lin et al., ). However these models are fundamentally limited to describing the kernel or lazy learning regime of neural networks (Chizat et al., ). Several works have found that this fails to capture the scaling laws of deep networks in the feature learning regime (Fort et al., ; Vyas et al., ; ; Bordelon et al., ). A theory of scaling laws that can capture consistent feature learning even in an infinite parameter $N\to\infty$ limit is especially pressing given the success of mean field and $\mu$-parameterizations which generate constant scale feature updates across model widths and depths (Mei et al., ; Geiger et al., ; Yang & Hu, ; Bordelon & Pehlevan, ; Yang et al., ; Bordelon et al., ; ). The training dynamics of the infinite width/depth limits in such models can significantly differ from the lazy training regime. Infinite limits which preserve feature learning are better descriptors of practical networks Vyas et al. (). Motivated by this, we ask the following:

Question: Under what conditions can feature learning improve the scaling law exponents of neural networks compared to lazy training regime?

### 1.1 Our Contributions

In this work, we develop a theoretical model of neural scaling laws that allows for improved scaling exponents compared to lazy training under certain settings. Our contributions are

- 1.
We propose a simple two-layer linear neural network model trained with a form of projected gradient descent. We show that this model reproduces power law scalings in training time, model size and training set size. The predicted scaling law exponents are summarized in terms of two parameters related to the data and architecture $(\alpha,\beta)$, which we refer to as the source and capacity exponents following Caponnetto & Vito ().
- 2.
We identify a condition on the difficulty of the learning task, measured by the source exponent $\beta$, under which feature learning can improve the scaling of the loss with time and with compute. For easy tasks, which we define as tasks with $\beta>1$ where the RKHS norm of the target is finite, there is no improvement in the power-law exponent while for hard tasks ($\beta<1$) that are outside the RKHS of the initial limiting kernel, there can be an improvement. For super-easy tasks $\beta>2-\frac{1}{\alpha}$, which have very low RKHS norm, variance from stochastic gradient descent (SGD) can alter the scaling law at large time. Figure 1 summarizes these results.
- 3.
We provide an approximate prediction of the compute optimal scaling laws for hard, easy tasks and super-easy tasks. Each of these regimes has a different exponent for the compute optimal neural scaling law. Table 1 summarizes these results.
- 4.
We test our predicted feature learning scalings by training deep nonlinear neural networks fitting nonlinear functions. In many cases, our predictions from the initial kernel spectra accurately capture the test loss of the network in the feature learning regime.

Figure: (a) $\lim_{N\to\infty}\mathcal{L}(t,N)\sim t^{-\chi(\beta)}$
Refer to caption: x1.png

### 1.2 Related Works

Our work builds on the recent results of Bordelon et al. () and Paquette et al. () which analyzed the SGD dynamics of a structured random feature model. Statics of this model have been analyzed by many prior works (Atanasov et al., ; Simon et al., ; Zavatone-Veth & Pehlevan, ; Bahri et al., ; Maloney et al., ). These kinds of models can accurately describe networks in the lazy learning regime. However, the empirical study of Vyas et al. () and some experiments in Bordelon et al. () indicate that the predicted compute optimal exponents were smaller than those measured in networks that learn features on real data. These latter works observed that networks train faster in the rich regime compared to lazy training. We directly address this gap in performance between lazy and feature learning neural networks by allowing the kernel features to adapt to the data. We revisit the computer vision settings of Bordelon et al. () and show that our new exponents more accurately capture the scaling law in the feature learning regime.

Other work has investigated when neural networks outperform kernels (Ghorbani et al., ). Ba et al. () and Abbe et al. () have shown how feature learning neural networks can learn low rank spikes in the hidden layer weights/kernels to help with sparse tasks while lazy networks cannot. Target functions with staircase properties, where learning simpler components aid learning of more complex components also exhibit significant improvements (with respect to a large input dimension) due to feature learning (Abbe et al., ; Dandi et al., ; Bardone & Goldt, ). Here, we consider a different setting. We ask whether feature learning can lead to improvements in power law exponents for the neural scaling law. The work of Paccolat et al. () asks a similar question in the case of a simple stripe model. Here we investigate whether the power law scaling exponent can be improved with feature learning in a model that only depends on properties of the initial kernel and the target function spectra.

## 2 Solvable Model of Scaling Laws With Feature Learning

We start by motivating and defining our model. Our goals is to build a simple model that exhibits feature learning in the infinite-width limit but also captures finite network size, finite batch SGD effects, and sample size effects that can significantly alter scaling behavior.

Following the notation of Bordelon et al. (), we introduce our model from the perspective of kernel regression. We first assume a randomly initialized neural network in an infinite-width limit where the neural tangent kernel (NTK) concentrates. We then diagonalize the initial infinite-width NTK. The resulting eigenfunctions $\bm{\psi}_{\infty}(\bm{x})\in\mathbb{R}^{M}$ have an inner product that define the infinite-width NTK $K_{\infty}(\bm{x},\bm{x}^{\prime})=\bm{\psi}_{\infty}(\bm{x})\cdot\bm{\psi}_{
\infty}(\bm{x}^{\prime})$. These eigenfunctions are orthogonal under the probability distribution of the data $p(\bm{x})$ with

$$ $\displaystyle\left<\bm{\psi}_{\infty}(\bm{x})\bm{\psi}_{\infty}(\bm{x})^{\top} \right>_{\bm{\bm{x}}\sim p(\bm{x})}=\bm{\Lambda}=\text{diag}(\lambda_{1},..., \lambda_{M}).$ (2) $$

We will often consider the case where $M\to\infty$ first so that these functions $\bm{\psi}_{\infty}(\bm{x})$ form a complete basis for the space of square integrable functions.

We next consider a finite sized model with $N$ parameters. We assume this model’s initial parameters are sampled from the same distribution as the infinite model and that the $N\to\infty$ limit recovers the same kernel $K_{\infty}(\bm{x},\bm{x}^{\prime})$. The finite $N$-parameter model, at initialization $t=0$, has $N$ eigenfeatures $\tilde{\bm{\psi}}(\bm{x},0)\in\mathbb{R}^{N}$. In the feature learning regime, these features will evolve during training.

The finite network’s learned function $f$ is expressed in terms of the lower dimensional features, while the target function $y(\bm{x})$ can be decomposed in terms of the limiting (and static) features $\bm{\psi}_{\infty}(\bm{x})$ with coefficients $\bm{w}^{\star}$. The instantaneous finite width features $\tilde{\bm{\psi}}(\bm{x},t)$ can also be expanded as a linear combination of the basis functions $\bm{\psi}_{\infty}(\bm{x})$ with coefficient matrix $\bm{A}(t)\in\mathbb{R}^{N\times M}$. We can therefore view our model as the following student-teacher setup

$$ $\displaystyle f(\bm{x},t)=\frac{1}{N}\bm{w}(t)\cdot$ $\displaystyle\tilde{\bm{\psi}}(\bm{x},t)\ ,\quad\ \tilde{\bm{\psi}}(\bm{x},t)= \bm{A}(t)\bm{\psi}_{\infty}(\bm{x})$ (3) $\displaystyle y(\bm{x})=\bm{w}^{\star}\cdot\bm{\psi}_{\infty}(\bm{x}).$ $$

If the matrix $\bm{A}(t)$ is random and static then gradient descent on this random feature model recovers the stochastic process analyzed by Bordelon et al. (); Paquette et al. (). In this work, we extend the analysis to cases where the matrix $\bm{A}(t)$ is also updated. We consider online training in the main text but discuss and analyze the case where samples are reused in Appendix D.

We allow $\bm{w}(t)$ to evolve with stochastic gradient descent (SGD) and $\bm{A}(t)$ evolve by projected SGD on a mean square error with batch size $B$. Letting $\bm{\Psi}_{\infty}(t)\in\mathbb{R}^{B\times M}$ represent a randomly sampled batch of $B$ points evaluated on the limiting features $\{\bm{\psi}_{\infty}(\bm{x}_{\mu})\}_{\mu=1}^{B}$ and $\eta$ to be the learning rate, our updates take the form

$$ $\displaystyle\bm{w}(t+1)-\bm{w}(t)=\eta\bm{A}(t)\left(\frac{1}{B}\bm{\Psi}_{ \infty}(t)^{\top}\bm{\Psi}_{\infty}(t)\right)\bm{v}^{0}(t)\ ,\ \bm{v}^{0}(t) \equiv\bm{w}_{\star}-\frac{1}{N}\bm{A}(t)^{\top}\bm{w}(t)$ $\displaystyle\bm{A}(t+1)-\bm{A}(t)=\eta\gamma\ \bm{w}(t)\bm{v}^{0}(t)^{\top} \left(\frac{1}{B}\bm{\Psi}_{\infty}(t)^{\top}\bm{\Psi}_{\infty}(t)\right)\left (\frac{1}{N}\bm{A}(0)^{\top}\bm{A}(0)\right).$ (4) $$

The fixed random projection $\left(\frac{1}{N}\bm{A}(0)^{\top}\bm{A}(0)\right)$ present in $\bm{A}(t)$’s dynamics ensure that the features cannot have complete access to the infinite width features $\bm{\psi}_{\infty}$ but only access to the initial $N$-dimensional features $\bm{A}(0)\bm{\psi}_{\infty}$. If this term were not present then there would be no finite parameter bottlenecks in the model and even a model with $N=1$ could fully fit the target function, leading to trivial parameter scaling laws(^1^11We could also solve this problem by training a model of the form $f=\bm{w}(t)^{\top}\bm{B}(t)\bm{A}\bm{\psi}$ where $\bm{w}(t)$ and $\bm{B}(t)$ are dynamical with initial condition $\bm{B}(0)=\bm{I}$ and $\bm{A}$ frozen and the matrix $\bm{B}(t)$ following gradient descent. We show that these two models are actually exhibit equivalent dynamics in Appendix B.). In this sense, the vector space spanned by the features $\tilde{\bm{\psi}}$ does not change over the course of training, but the finite-width kernel Hilbert space does change its kernel: $\tilde{\bm{\psi}}(\bm{x},t)\cdot\tilde{\bm{\psi}}(\bm{x}^{\prime},t)$. Feature learning in this space amounts to reweighing the norms of the existing features. We have chosen $\bm{A}(t)$ to have dynamics similar to the first layer weight matrix of a linear neural network. As we will see, this is enough to lead to an improved scaling exponent.

The hyperparameter $\gamma$ sets the timescale of $\bm{A}$’s dynamics and thus controls the rate of feature evolution. The $\gamma\to 0$ limit represents the lazy learning limit Chizat et al. () where features are static and coincides with a random feature model dynamics of Bordelon et al. (); Paquette et al. (). The test error after $t$ steps on a $N$ parameter model with batch size $B$ is

$$ $\displaystyle\mathcal{L}(t,N,B,\gamma)$ $\displaystyle\equiv\left<\left[\bm{\psi}_{\infty}(\bm{x})\cdot\bm{w}^{*}- \tilde{\bm{\psi}}(\bm{x},t)\cdot\bm{w}(t)\right]^{2}\right>_{\bm{x}\sim p(\bm{ x})}=\bm{v}^{0}(t)^{\top}\bm{\Lambda}\bm{v}^{0}(t).$ (5) $$

In the next sections we will work out a theoretical description of this model as a function of the spectrum $\bm{\Lambda}$ and the target coefficients $\bm{w}^{\star}$. We will then specialize to power-law spectra and target weights and study the resulting scaling laws.

## 3 Dynamical Mean Field Theory of the Model

We can consider the dynamics for random $\bm{A}(0)$ and random draws of data during SGD in the limit of $M\to\infty$ and $N,B\gg 1$(^2^22Alternatively, we can operate in a proportional limit with $N/M,B/M$ approaching constants. This version of the theory would be exact with no finite size fluctuations.). This dimension-free theory is especially appropriate for realistic trace class kernels where $\left<K_{\infty}(\bm{x},\bm{x}^{\prime})\right>_{\bm{x}}=\sum_{k}\lambda_{k}<\infty$ (equivalent to $\alpha>1$), which is our focus. Define $w_{k}^{\star},v_{k}(t)$ to be respectively the components of $\bm{w}^{\star},\bm{v}^{0}(t)$ in the $k$th eigenspace of $\bm{\Lambda}$. The error variables $v^{0}_{k}(t)$ are given by a stochastic process, and yield deterministic prediction for the loss $\mathcal{L}(t,N,B,\gamma)$, analogous to the results of Bordelon et al. ().

Since the resulting dynamics for $v^{0}_{k}(t)$ at $\gamma>0$ are nonlinear and cannot be expressed in terms of a matrix resolvent, we utilize dynamical mean field theory (DMFT), a flexible approach for handling nonlinear dynamical systems driven by random matrices (Sompolinsky & Zippelius, ; Helias & Dahmen, ; Mannelli et al., ; Mignacco et al., ; Gerbelot et al., ; Bordelon et al., ). Most importantly, the theory gives closed analytical predictions for $\mathcal{L}(t,N,B,\gamma)$. We defer the derivation and full DMFT equations to the Appendix C. The full set of closed DMFT equations are given in Equation equation 26 for online SGD and Equation equation 41 for offline training with data repetition. Informally, this DMFT computes a closed set of equations for the correlation and response functions for a collection of time-varying vectors $\mathcal{V}=\{\bm{v}^{0}(t),\bm{v}^{1}(t),\bm{v}^{2}(t),\bm{v}^{3}(t),\bm{v}^{
4}(t)\}_{t\in\{0,1,...\}}$ including,

$$ $\displaystyle C_{0}(t,s)=\bm{v}^{0}(t)^{\top}\bm{\Lambda}\bm{v}^{0}(s)\ ,\ C_{ 3}(t,s)=\frac{1}{N}\bm{v}^{3}(t)\cdot\bm{v}^{3}(s)\ ,\ R_{0,2}(t,s)=\text{Tr} \bm{\Lambda}\left[\frac{\delta\bm{v}^{0}(t)}{\delta\bm{v}^{2}(s)^{\top}}\right],$ $$

where $\frac{\delta\bm{v}^{k}(t)}{\delta\bm{v}^{\ell}(s)^{\top}}$ represents the response of $\bm{v}^{k}$ at time $t$ to a small perturbation to $\bm{v}^{\ell}$ at time $s$. The test loss can be computed from $\mathcal{L}(t)=C_{0}(t,t)$. This theory is derived generally for any spectrum $\lambda_{k}$ and any target $w^{\star}_{k}$. In the coming sections we will examine approximate scaling behavior of the loss when the spectrum follows a power law.

## 4 Power Law Scalings from Power Law Features

We consider initial kernels that satisfy source and capacity conditions as in (Caponnetto & Vito, ; Pillaud-Vivien et al., ; Cui et al., ; ). These conditions measure the rate of decay of the spectrum of the initial infinite width kernel $K_{\infty}(x,x^{\prime})$ and target function $y(x)$ in that basis. Concretely, we consider settings with the following power law scalings:

$$ $\displaystyle\lambda_{k}\sim k^{-\alpha}\ ,\quad\ \sum_{\ell>k}\lambda_{\ell}( w_{\ell}^{*})^{2}\sim\lambda_{k}^{-\beta}\sim k^{-\alpha\beta}.$ (6) $$

The exponent $\alpha$ is called the capacity and measures the rate of decay of the initial kernel eigenvalues. We will assume this exponent is greater than unity $\alpha>1$ since the limiting $N\to\infty$ kernel should be trace class. The exponent $\beta$ is called the source and quantifies the difficulty of the task under kernel regression with $K_{\infty}$.(^3^33The source exponent $r$ used in (Pillaud-Vivien et al., ) and other works is given by $2r=\beta$.) The RKHS norm $|\cdot|_{\mathcal{H}}^{2}$ of the target function is given by:

$$ $\displaystyle|y|_{\mathcal{H}}^{2}=\sum_{k}(w^{\star}_{k})^{2}=\sum_{k}k^{- \alpha(\beta-1)-1}\approx\begin{cases}\frac{1}{\alpha(\beta-1)}&\beta>1\\ \infty&\beta<1.\end{cases}$ (7) $$

While the case of finite RKHS norm ($\beta>1$) is often assumed in analyses of kernel methods that rely on norm-based bounds, such as (Bartlett & Mendelson, ; Bach, ), the $\beta<1$ case is actually more representative of real datasets. This was pointed out in (Wei et al., ). This can be seen by spectral diagonalizations performed on real datasets in (Bahri et al., ; Bordelon et al., ) as well as in experiments in Section 5.2. We stress this point since the behavior of feature learning with $\beta>1$ and $\beta<1$ will be strikingly different in our model.

##### General Scaling Law in the Lazy Limit

For the purposes of deriving compute optimal scaling laws, the works of Bordelon et al. () and Paquette et al. () derived precise asymptotics for the loss curves under SGD. For the purposes of deriving compute optimal scaling laws, these asymptotics can be approximated as the following sum of power laws at large $t,N$

$$ $\displaystyle\lim_{\gamma\to 0}\mathcal{L}(t,N,B,\gamma)\approx\underbrace{t^{ -\beta}}_{\text{Limiting Gradient Flow}}+\underbrace{N^{-\alpha\min\{2,\beta\} }}_{\text{Model Bottleneck}}+\underbrace{\frac{1}{N}t^{-(1-\frac{1}{\alpha})}} _{\text{Finite $N$ Transient}}+\underbrace{\frac{\eta}{B}t^{-(2-\frac{1}{ \alpha})}}_{\text{SGD Transient}}.$ (8) $$

The first terms represent bottleneck/resolution-limited scalings in the sense that taking all other quantities to infinity and studying scaling with the last (Bahri et al., ). The first term gives the loss dynamics of population gradient flow ($N,B\to\infty$) while the second (model bottleneck) term describes $t\to\infty$ limit of the loss which depends on $N$. The third and fourth terms are mixed transients that arise from the perturbative finite model and batch size effects. While Bordelon et al. () focused on hard tasks where $\beta<1$ where the first two terms dominate when considering compute optimal scaling laws, Paquette et al. () also discussed two other phases of the easy task regime where the final two terms can become relevant for the compute optimal scaling.

##### General Scaling Law in the Feature Learning Regime

For $\gamma>0$, approximations to our precise DMFT equations under power law spectra give the following sum of power laws

$$ $\displaystyle\mathcal{L}(t,N,B,\gamma)\approx\underbrace{t^{-\beta\max\left\{1 ,\frac{2}{1+\beta}\right\}}}_{\text{Limiting Gradient Flow}}+\underbrace{N^{- \alpha\min\{2,\beta\}}}_{\text{Model Bottleneck}}+\underbrace{\frac{1}{N}t^{-( 1-\frac{1}{\alpha})\max\left\{1,\frac{2}{1+\beta}\right\}}}_{\text{Finite $N$ Transient}}+\underbrace{\frac{\eta}{B}t^{-(2-\frac{1}{\alpha})\max\left\{1, \frac{2}{1+\beta}\right\}}}_{\text{SGD Transient}}.$ (9) $$

We see that in the rich regime, all exponents except for the model bottleneck are either the same or are improved. For easy tasks and super-easy tasks where $\beta>1$, we recover the same approximate scaling laws as those computed in the linear model of Bordelon et al. () and Paquette et al. (). For hard tasks, $\beta<1$, all exponents except for the model bottleneck term are improved. Below we will explain why each of these terms can experience an improvement in the $\beta<1$ case. We exhibit a phase diagram all of the cases highlighted in equation 8, equation 9 in Figure 1.

##### Accelerated Training in Rich Regime

The key distinction between our model and the random feature model ($\gamma=0$) is the limiting gradient flow dynamics, which allow for acceleration due to feature learning. For nonzero feature learning $\gamma>0$, our theory predicts that in the $N\to\infty$ limit, the loss scales as a power law $\mathcal{L}(t)\sim t^{-\chi}$ where the exponent $\chi$ satisfies the following self-consistent equation

$$ $\displaystyle\chi(\beta)=-\lim_{t\to\infty}\frac{1}{\ln t}\ln\left[\sum_{k}(w^ {\star}_{k})^{2}\lambda_{k}\exp\left(-\lambda_{k}\left[t+\gamma t^{2-\chi} \right]\right)\right]=\beta\max\left\{1,\frac{2}{1+\beta}\right\}.$ (10) $$

We derive this equation in Appendix E.
We see that if $\beta>1$ then we have the same scaling law as a lazy learning model $\mathcal{L}(t)\sim t^{-\beta}$. However, if the task is sufficiently hard ($\beta<1$), then the exponent is increased to $\chi=\frac{2\beta}{1+\beta}>\beta$.

This acceleration is caused by the fact that the effective dynamical kernel $K(t)$ defined by the dynamics of our features $\tilde{\bm{\psi}}(\bm{x},t)$ diverges as a powerlaw $K(t)\sim t^{1-\chi}$ when $\beta<1$ (see Appendix E). As a consequence, at time $t$, the model is learning mode $k_{\star}(t)\sim t^{(2-\chi)/\alpha}$ which gives a loss

$$ $\displaystyle\mathcal{L}(t)\sim\sum_{k>k_{\star}}(w^{\star}_{k})^{2}\lambda_{k }\sim t^{-\beta(2-\chi)}=t^{-\beta\max\left\{1,\frac{2}{1+\beta}\right\}}.$ (11) $$

While our model predicts that the scaling exponent only changes for hard tasks where $\beta<1$, it also predicts an overall decrease in training loss as $\gamma$ increases for either easy or hard tasks. In Figure 2 (a)-(b) we show the the $N,B\to\infty$ limit of our theory at varying values of $\gamma$. For easy tasks $\beta>1$, the models will always follow $\mathcal{L}\sim t^{-\beta}$ at late time, but with a potentially reduced constant when $\gamma$ is large. For hard tasks (Fig. 2 (b)) the scaling exponent improves $\mathcal{L}\sim t^{-\frac{2\beta}{1+\beta}}$ for $\gamma>0$.

Figure: (a) Easy Task $\beta=1.2$ with $N,B\to\infty$
Refer to caption: x4.png

##### Model Bottleneck Scalings

Our theory can be used to compute finite $N$ effects in the rich regime during SGD training. In this case, the dynamics smoothly transition between following the gradient descent trajectory at early time to an asymptote that depends on $N$ as $t\to\infty$. In Figure 2 (c)-(d) we illustrate these learning curves from our theory and from finite $N$ simulations, showing a good match of the theory to experiment.

We derive the asymptotic scaling of $N^{-\alpha\min\{2,\beta\}}$ in Appendix E.3. Intuitively, at finite $N$, the dynamics only depend on the filtered signal $\left(\frac{1}{N}\bm{A}(0)^{\top}\bm{A}(0)\right)\bm{w}_{\star}$. Thus the algorithm can only estimate, at best, the top $N$ components of $\bm{w}_{\star}$, resulting in the following $t\to\infty$ loss

$$ $\displaystyle\mathcal{L}(N)\sim\sum_{k>N}k^{-\alpha\beta-1}\sim N^{-\alpha \beta}.$ (12) $$

##### SGD Noise Effects

The variance in the learned model predictions due to random sampling of minibatches during SGD also alters the mean field prediction of the loss. In Figure 3, we show SGD noise effects from finite batch size $B$ for hard $\beta<1$ and super easy $\beta>2-1/\alpha$ tasks.

Figure: (a) SGD Transients Hard $\beta=0.5$
Refer to caption: x8.png

##### Compute Optimal Scaling Laws in Feature Learning Regime

At a fixed compute budget $C=Nt$, one can determine how to allocate compute towards training time $t$ and model size $N$ using our derived exponents from the previous sections. Choosing $N,t$ optimally, we derive the following compute optimal scaling laws $\mathcal{L}_{\star}(C)$ in the feature learning regime $\gamma>0$. These are also summarized in Figure 1. (^4^44The three regimes of interest correspond to Phases I,II,III in Paquette et al. (). These are the only relevant regimes for trace-class $\left<K_{\infty}(\bm{x},\bm{x})\right>_{\bm{x}}=\sum_{k}\lambda_{k}<\infty$ (finite variance) kernels (equivalent to $\alpha>1$).)

- 1.
Hard task regime ($\beta<1)$: the compute optimum balances the population gradient flow term $t^{-\frac{2\beta}{1+\beta}}$ and the model bottleneck $N^{-\alpha\beta}$.
- 2.
Easy tasks ($1<\beta<2-\frac{1}{\alpha}$): the compute optimum compares gradient flow term $t^{-\beta}$ to finite $N$ transient term $\frac{1}{N}t^{-1+1/\alpha}$
- 3.
Super easy tasks ($\beta>2-\frac{1}{\alpha}$): compute optimum balances the finite N transient $\frac{1}{N}t^{-1+1/\alpha}$ and SGD transient terms $\frac{1}{B}t^{-2+\frac{1}{\alpha}}$.

We work out the complete compute optimal scaling laws for these three settings by imposing the constraint $C=Nt$, identifying the optimal choice of $N$ and $t$ at fixed $t$ and verifying the assumed dominant balance. We summarize the three possible compute scaling exponents in Table 1.

**Table 1: Compute optimal scaling exponents $r_{C}$ for the loss $\mathcal{L}_{\star}(C)\sim C^{-r_{C}}$ for tasks of varying difficulty in the feature learning regime. For $\beta>1$, the exponents coincide with the lazy model analyzed by Bordelon et al. (); Paquette et al. (), while for hard tasks they are improved.**
| Task Difficulty | Hard $\beta<1$ | Easy $1<\beta<2-1/\alpha$ | Super-Easy $\beta>2-1/\alpha$ |
| --- | --- | --- | --- |
| Lazy ($\gamma=0$) | $\frac{\alpha\beta}{\alpha+1}$ | $\frac{\alpha\beta}{\alpha\beta+1}$ | $1-\frac{1}{2\alpha}$ |
| Rich ($\gamma>0$) | $\frac{2\alpha\beta}{\alpha(1+\beta)+2}$ | $\frac{\alpha\beta}{\alpha\beta+1}$ | $1-\frac{1}{2\alpha}$ |

In Figure 4 we compare the compute optimal scaling laws in the hard and easy regimes. We show that the predicted exponents are accurate. In Figure 3 we illustrate the influence of SGD noise on the learning curve in the super easy regime and demonstrate that the large $C$ compute optimal scaling law is given by $C^{-1+\frac{1}{2\alpha}}$.

Figure: (a) Hard Task Scaling $\beta=0.5$
Refer to caption: x10.png

## 5 Experiments with Deep Nonlinear Neural Networks

While our theory accurately describes simulations of our solvable model, we now aim to test if these new exponents are predictive when training deep nonlinear neural networks.

### 5.1 Sobolev Spaces on the Circle

We first consider training multilayer nonlinear MLPs with nonlinear activation function $\phi(h)=\left[\text{ReLU}(h)\right]^{q_{\phi}}$ in the mean field parameterization/$\mu$P (Geiger et al., ; Yang & Hu, ; Bordelon & Pehlevan, ) with dimensionless (width-independent) feature learning parameter $\gamma_{0}$. We consider fitting target functions $y(x)$ with $x=[\cos(\theta),\sin(\theta)]^{\top}$ on the circle $\theta\in[0,2\pi]$. The eigenfunctions for randomly initialized infinite width networks are the Fourier harmonics. We consider target functions $y(\theta)$ with power-law Fourier spectra while the kernels at initialization $K(\theta,\theta^{\prime})$ also admit a Fourier eigenexpansion

$$ $\displaystyle y(\theta)=\sum_{k=1}^{\infty}k^{-q}\cos(k\theta)\ ,\ K(\theta, \theta^{\prime})=\sum_{k=1}^{\infty}\lambda_{k}\cos(k(\theta-\theta^{\prime})).$ (13) $$

We show that the eigenvalues of the kernel at initialization decay as $\lambda_{k}\sim k^{-2q_{\phi}}$ in the Appendix A. The capacity and source exponents $\alpha,\beta$ required for our theory can be computed from $q$ and $q_{\phi}$ as

$$ $\displaystyle\alpha=2q_{\phi}\ ,\ \beta=\frac{2q-1}{2q_{\phi}}$ (14) $$

Thus task difficulty can be manipulated by altering the target function or the nonlinear activation function of the neural network. We show in Figure 5 examples of online training in this kind of network on tasks and architectures of varying $\beta$. In all cases, our theoretical prediction of $t^{-\beta\max\{1,\frac{2}{1+\beta}\}}$ provides a very accurate prediction of the scaling law.

Figure: (a) ReLU $(q_{\phi}=1.0)$ with varying $\beta=\frac{2q-1}{2q_{\phi}}$
Refer to caption: x13.png

### 5.2 Computer Vision Tasks (MNIST and CIFAR)

We next study networks trained on MNIST and CIFAR image recognition tasks. Our motivation is to study networks training in the online setting over several orders of magnitude in time. To this end, we adopt larger versions of these datasets: “MNIST-1M” and CIAFR-5M. We generate MNIST-1M using the denoising diffusion model (Ho et al., ) in Pearce (). We use CIFAR-5M from Nakkiran et al. (). Earlier results in Refinetti et al. () show that networks trained on CIFAR-5M have very similar trajectories to those trained on CIFAR-10 without repetition. The resulting scaling plots are provided in Figure 6. MNIST-1M scaling is very well captured by the our theoretical scaling exponents. The CIFAR-5M scaling law exponent at large $\gamma_{0}$ first follows our predictions, but later enters a regime with a scaling exponent larger than what our theoretical model predicts.

Figure: (a) CNNs on MNIST-1M, $\beta=0.30$
Refer to caption: x15.png

## 6 Discussion

We proposed a simple model of learning curves in the rich regime where the original features can evolve as a linear combination of the initial features. While the theory can give a quantitatively accurate for the online learning scaling exponent in the rich regime for hard tasks, the CIFAR-5M experiment suggests that additional effects in nonlinear networks can occur after sufficient training. However, there are many weaker predictions of our theory that we suspect to hold in a wider set of settings, which we enumerate below.

##### Source Hypothesis:

Feature Learning Only Improves Scaling For $\beta<1$. Our model makes a general prediction that feature learning does not improve the scaling laws for tasks within the RKHS of the initial infinite width kernel. Our experiments with ReLU networks fitting functions in different Sobolev spaces with $\beta>1$ support this hypothesis. Since many tasks using real data appear to fall outside the RKHS of the initial infinite width kernels, this hypothesis suggests that lazy learning would not be adequate to describe neural scaling laws on real data, consistent with empirical findings (Vyas et al., ).

##### Insignificance of SGD for Hard Tasks

Recent empirical work has found that SGD noise has little impact in online training of deep learning models (Vyas et al., ; Zhao et al., ). Our theory suggests this may be due to the fact that SGD transients are always suppressed for realistic tasks which are often outside the RKHS of the initial kernel. The regime in which feature learning can improve the scaling law in our model is precisely the regime where SGD transients have no impact on the scaling behavior.

##### Ordering of Models in Lazy Limit Preserved in Feature Learning Regime

An additional interesting prediction of our theory is that the ordering of models by performance in the lazy regime is preserved is the same as the ordering of models in the feature learning regime. If model A outperforms model B on a task in the lazy limit ($\beta_{A}>\beta_{B}$), then model A will also perform better in the rich regime $\chi(\beta_{A})>\chi(\beta_{B})$ (see Figure 1). This suggests using kernel limits of neural architectures for fast initial architecture search may be viable, despite failing to capture feature learning (Park et al., ). This prediction deserves a greater degree of stress testing.

##### Limitations and Future Directions

There are many limitations to the current theory. First, we study mean square error loss with SGD updates, while most modern models are trained on cross-entropy loss with adaptive optimizers (Everett et al., ). Understanding the effect of adaptive optimizers or preconditioned updates on the scaling laws represents an important future direction. In addition, our model treats the learned features as linear combinations of the initial features, an assumption which may be violated in finite width neural networks. Lastly, while our theory is very descriptive of nonlinear networks on several tasks, we did identify a noticeable disagreement on CIFAR-5M after sufficient amounts of training. Versions of our model where the learned features are not within the span of the initial features or where the matrix $\bm{A}$ undergoes different dynamics may provide a promising avenue of future research to derive effective models of neural scaling laws.

### Acknowledgements

We would like to thank Jacob Zavatone-Veth, Jascha Sohl-Dickstein, Courtney Paquette, Bruno Loureiro, and Yasaman Bahri for useful discussions.

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.

## Appendix

## Appendix A Additional Experiments

Figure: (a) Kernel eigenspectra $\alpha=2q_{\phi}$
Refer to caption: x17.png

Figure: Figure 8: A depth $L=4$ decoder-only transformer ($16$ heads with $d_{\text{head}}=128$) trained on next-word prediction with SGD on the C4 dataset tokenized with the SentencePiece tokenizer. We plot cross entropy loss and fit a powerlaw $t^{-\beta}$ to lazy learning curve $\gamma=10^{-4}$ over the interval from $10^{6}$ to $3\times 10^{9}$ tokens. We then compute the new predicted exponent $t^{-\beta(2-\beta)}$ and compare to a simulation at $\gamma=0.1$. Though our theoretical prediction of a doubling of the scaling exponent was derived in the context of MSE, the new scaling exponent fits the data somewhat well for this setting at early times.
Refer to caption: x19.png

## Appendix B Further Discussion of Models

Figure: Figure 9: Three different models studied in this work and prior work. Black weights are frozen while orange weights are trainable. a) A linear random feature model with only the readout weights trainable. This model was studied in Maloney et al. (); Bordelon et al. (); Paquette et al. () as a solvable model of neural scaling laws. b) A two layer linear network with both weights trainable. This model does not incur a bottleneck due to finite width but undergoes feature learning, which improves the scaling of the loss with time. We study pure linear neural networks in Appendix G. In the main text, we train this model with a projected version of gradient descent. This is equivalent to c) and gives both finite-parameter bottlenecks as well as improvements to scaling due to feature learning.
Refer to caption: x20.png

We seek a model that incorporates both the bottle-necking effects of finite width observed in recent linear random feature models of scaling laws while still allowing for a notion of feature learning. Exactly solvable models of feature learning networks are relatively rare. Here, we take inspiration from the linear neural network literature Saxe et al. (). Linear neural networks exhibit both lazy and rich regimes of learning Woodworth et al. (), in which they can learn useful task-relevant features in a way that can be analytically studied Atanasov et al. (). In our work, we go beyond these models of linear neural networks and show that linear neural networks trained on data under source and capacity conditions can improve the convergence rate compared to that predicted by kernel theory.

The model introduced in section 2 is give by a two-layer linear network acting on the $\bm{\psi}_{\infty}(\bm{x})$:

$$ $f(\bm{x})=\frac{1}{N}\bm{w}^{\top}\bm{A}\bm{\psi}_{\infty}(\bm{x}).$ (15) $$

There, we constrained it to update its weights by a form of projected gradient descent as given by Equation 2. Here, we show that running this projected gradient descent is equivalent to running ordinary gradient descent on a two layer linear network after passing $\bm{\psi}_{\infty}$ through random projections. Define $\bm{A}_{0}=\bm{A}(0)$ and $\bm{B}(t)=\bm{A}(t)\bm{A}_{0}^{+}$ where $+$ denotes the Moore-Penrose psuedoinverse. Then assuming $N<M$ and $\bm{A}_{0}$ is rank $N$, we have

$$ $\bm{B}(t)\bm{A}_{0}=\bm{A}(t),\quad\bm{B}(0)=\mathbf{I}_{N\times N}.$ (16) $$

Now consider taking $\bm{\psi}_{\infty}$ and passing it through $\bm{A}_{0}$, and then training $\bm{B}$, $\bm{w}$ with ordinary gradient descent. We have update equations:

$$ $\displaystyle\bm{w}(t+1)-\bm{w}(t)$ $\displaystyle=\eta\bm{B}(t)\bm{A}_{0}\left(\frac{1}{B}\bm{\Psi}_{\infty}^{\top }\bm{\Psi}_{\infty}\right)\bm{v}^{0}(t)$ (17) $\displaystyle\bm{B}(t+1)-\bm{B}(t)$ $\displaystyle=\eta\bm{w}(t)\bm{v}^{0}(t)^{\top}\left(\frac{1}{B}\bm{\Psi}_{ \infty}^{\top}\bm{\Psi}_{\infty}\right)\frac{1}{N}\bm{A}_{0}^{\top}$ $$

Multiplying the second equation by $\bm{A}_{0}$ on the right recovers equation 2. Here, $\gamma$ acts as a rescaling of the learning rate for the $\bm{B}$ update equations. We illustrate this model, as well as the linear random feature and linear neural network models in Figure 9.

Several papers Maloney et al. (); Atanasov et al. (); Bordelon et al. (); Atanasov et al. () have studied the model given in equation 15 with frozen $\bm{A}$ under the following interpretation. The samples $\bm{\psi}\in\mathbb{R}^{D}$ correspond to the dataset as expressed in the space of an infinitely wide NTK at initialization. These are passed through a set of frozen random weights $\bm{W}_{1}$, which are thought of as the projection from the infinite width network to the finite-width empirical NTK, corresponding to a lazy network. From there, the final layer weights are not frozen and perform the analog of regression with the finite-width NTK. In Bordelon et al. (), this model was shown to reproduce many of the compute optimal scaling laws observed in practice. It was also however shown there that the scaling laws for lazy networks are very different from those observed for feature-learning networks.

Our motivation is to develop a simple and solvable model of how the finite-width network features $\tilde{\bm{\psi}}(\bm{x};t)=\bm{A}(t){\bm{\psi}}_{\infty}(\bm{x})$ might evolve to learn useful features. The projected linear model defined above states that the $\tilde{\bm{\psi}}$ recombine themselves in such a way so that the empirical neural tangent kernel $\tilde{\bm{\psi}}(\bm{x};t)\cdot\tilde{\bm{\psi}}(\bm{x}^{\prime};t)$ is better aligned to the task. The simple model of a linear neural network is rich enough to yield an improved power law, while still being analytically tractable.

## Appendix C Derivation of the Mean Field Equations

In this setting, we derive the mean field equations for the typical test loss $\mathcal{L}(t,N,B)$ as a function of training time. To accomplish this, we have to perform disorder averages over the random matrices $\bm{A}(0)$ and $\{\bm{\Psi}(t)\}_{t=0}^{\infty}$. We start by defining the following collection of fields

$$ $\displaystyle\bm{v}^{0}(t)=\bm{w}^{\star}-\frac{1}{N}\bm{A}(t)^{\top}\bm{w}(t)$ $\displaystyle\bm{v}^{1}(t)=\bm{\Psi}(t)\bm{v}^{0}(t)\ ,\ \bm{v}^{2}(t)=\frac{1 }{B}\bm{\Psi}(t)^{\top}\bm{v}^{1}(t)$ $\displaystyle\bm{v}^{3}(t)=\bm{A}(0)\bm{v}^{2}(t)\ ,\ \bm{v}^{4}(t)=\frac{1}{N }\bm{A}(0)^{\top}\bm{v}^{3}(t)$ $\displaystyle\bm{v}^{w}(t)=\frac{1}{N}\bm{A}(0)^{\top}\bm{w}(t)$ (18) $$

From these primitive fields, we can simplify the dynamics of $\bm{A},\bm{w}(t)$

$$ $\displaystyle\bm{A}(t)$ $\displaystyle=\bm{A}(0)+\eta\gamma\sum_{s<t}\bm{w}(s)\bm{v}^{4}(s)^{\top}$ $\displaystyle\bm{w}(t+1)$ $\displaystyle=\bm{w}(t)+\eta\bm{A}(t)\bm{v}^{2}(t)$ $\displaystyle=\bm{w}(t)+\eta\bm{v}^{3}(t)+\eta^{2}\gamma\sum_{s<t}\bm{w}(s) \left[\bm{v}^{4}(s)\cdot\bm{v}^{2}(t)\right]$ $\displaystyle=\bm{w}(t)+\eta\bm{v}^{3}(t)+\eta^{2}\gamma\sum_{s<t}\bm{w}(s)C_{ 3}(t,s)$ (19) $$

where we introduced the correlation function $C_{3}(t,s)\equiv\frac{1}{N}\bm{v}^{3}(t)\cdot\bm{v}^{3}(s)=\bm{v}^{2}(t)\cdot
\bm{v}^{4}(s)$. Similarly for $\bm{v}^{0}(t)$ and $\bm{v}^{w}(t)$ we have

$$ $\displaystyle\bm{v}^{0}(t)$ $\displaystyle=\bm{w}^{\star}-\bm{v}^{w}(t)-\eta\gamma\sum_{s<t}\bm{v}^{4}(s)C_ {w}(t,s)$ $\displaystyle\bm{v}^{w}(t+1)$ $\displaystyle=\bm{v}^{w}(t)+\eta\bm{v}^{4}(t)+\eta^{2}\gamma\sum_{s<t}\bm{v}^{ w}(t)C_{3}(t,s)$ (20) $$

where we introduced $C_{w}(t,s)\equiv\frac{1}{N}\bm{w}(t)\cdot\bm{w}(s)$. We see that, conditional on the correlation function $C_{3}(t,s)$, the vector $\bm{v}^{w}(t)$ can be interpreted as a linear filtered version of $\{\bm{v}^{4}(s)\}_{s<t}$ and is thus redundant. In addition, we no longer have to work with the random matrix $\bm{A}(t)$ but can rather track projections of this matrix on vectors of interest. Since all dynamics only depend on the random variables $\mathcal{V}=\{\bm{v}^{0},\bm{v}^{1},\bm{v}^{2},\bm{v}^{3}\}$, we therefore only need to characterize the joint distribution of these variables.

##### Disorder Averages

We now consider the averages over the random matrices which appear in the dynamics $\{\bm{\Psi}(t)\}_{t\in\mathbb{N}}$ and $\bm{A}(0)$. This can be performed with either a path integral or a cavity derivation following the techniques of Bordelon et al. (). After averaging over $\{\bm{\Psi}(t)\}_{t=0}^{\infty}$, one obtains the following process for $v^{1}(t)$ and $v^{2}_{k}(t)$

$$ $\displaystyle v^{1}(t)=u^{1}(t)\ ,\ u^{1}(t)\sim\mathcal{N}(0,\delta(t-s)C_{0} (t,t))$ $\displaystyle v^{2}_{k}(t)=u^{2}_{k}(t)+\lambda_{k}v^{0}_{k}(t)\ ,\ u^{2}_{k}( t)\sim\mathcal{N}\left(0,B^{-1}\delta(t-s)\lambda_{k}C_{1}(t,t)\right).$ (21) $$

where the correlation functions $C_{0}$ and $C_{1}$ have the forms

$$ $\displaystyle C_{0}(t,s)=\sum_{k}\lambda_{k}\left<v^{0}_{k}(t)v^{0}_{k}(s) \right>\ ,\ C_{1}(t,s)=\left<v^{1}(t)v^{1}(s)\right>$ (22) $$

The average over the matrix $\bm{A}(0)$ couples the dynamics for $\bm{v}^{3}(t),\bm{v}^{4}(t)$ resulting in the following

$$ $\displaystyle v^{3}(t)=u^{3}(t)+\frac{1}{N}\sum_{s<t}R_{2,4}(t,s)v^{3}(s)\ ,\ u^{3}(t)\sim\mathcal{N}(0,C_{2}(t,s))$ $\displaystyle v^{4}_{k}(t)=u^{4}_{k}(t)+\sum_{s<t}R_{3}(t,s)v^{2}_{k}(s)\ ,\ u ^{4}_{k}\sim\mathcal{N}\left(0,N^{-1}C_{3}(t,s)\right)$ (23) $$

where

$$ $\displaystyle C_{2}(t,s)=\sum_{k}\left<v^{2}_{k}(t)v^{2}_{k}(s)\right>\ ,\ C_{ 3}(t,s)=\left<v^{3}(t)v^{3}(s)\right>$ (24) $$

Lastly, we have the following single site equation for $w(t)$ which can be used to compute $C_{w}(t,s)$

$$ $\displaystyle w(t+1)-w(t)=\eta v^{3}(t)+\eta^{2}\gamma\sum_{s<t}C_{3}(t,s)w(s).$ (25) $$

##### Final DMFT Equations for our Model for Online SGD

The complete governing equations for the test loss evolution after averaging over the random matrices can be obtained from the following stochastic processes which are driven by Gaussian noise sources $\{u^{2}_{k}(t),u^{3}(t),u^{4}_{k}(t)\}$. Letting $\left<\cdot\right>$ represent averages over these sources of noise, the equations close as

$$ $\displaystyle v^{0}_{k}(t)=w^{\star}_{k}-v^{w}_{k}(t)-\eta\gamma\sum_{s<t}C_{w }(t,s)v^{4}_{k}(s)$ $\displaystyle v^{w}_{k}(t+1)=v^{w}_{k}(t)+\eta v^{4}_{k}(t)+\eta^{2}\gamma\sum _{s<t}C_{3}(t,s)v^{w}_{k}(s)$ $\displaystyle v^{2}_{k}(t)=u^{2}_{k}(t)+\lambda_{k}v^{0}_{k}(t)\ ,\ u^{2}_{k}( t)\sim\mathcal{N}\left(0,B^{-1}\lambda_{k}\delta(t-s)C_{0}(t,t)\right)$ $\displaystyle v^{3}(t)=u^{3}(t)+\frac{1}{N}\sum_{s<t}R_{2,4}(t,s)v^{3}(s)\ ,\ u^{3}(t)\sim\mathcal{N}(0,C_{2}(t,s))$ $\displaystyle w(t+1)=w(t)+\eta v^{3}(t)+\eta^{2}\gamma\sum_{s<t}C_{3}(t,s)w(s)$ $\displaystyle v^{4}_{k}(t)=u^{4}_{k}(t)+\sum_{s<t}R_{3}(t,s)v^{2}_{k}(s)\ ,\ u ^{4}_{k}(t)\sim\mathcal{N}(0,N^{-1}C_{3}(t,s))$ $\displaystyle R_{2,4}(t,s)=\sum_{k}\left<\frac{\partial v^{2}_{k}(t)}{\partial u ^{4}_{k}(s)}\right>\ ,\ R_{3}(t,s)=\left<\frac{\partial v^{3}(t)}{\partial u^{ 3}(s)}\right>$ $\displaystyle C_{0}(t,s)=\sum_{k}\lambda_{k}\left<v^{0}_{k}(t)v^{0}_{k}(s) \right>\ ,\ C_{2}(t,s)=\sum_{k}\left<v^{2}_{k}(t)v^{2}_{k}(s)\right>$ (26a) $$

##### Closing the DMFT Correlation and Response as Time × \times × Time Matrices

We can try using these equations to write a closed form expression for $v^{0}_{k}(t)$ which determines the generalization error. First, we start by solving the equations for $\bm{v}^{w}_{k}=\text{Vec}\{v^{w}_{k}(t)\}_{t\in\mathbb{N}}$. We introduce a step function matrix which is just lower triangular matrix $[\bm{\Theta}]_{t,s}=\eta\Theta(t-s)$

$$ $\displaystyle\bm{v}^{w}_{k}=\left[\bm{I}-\eta\gamma\bm{\Theta}\bm{C}_{3}\right ]^{-1}\bm{\Theta}\bm{v}^{4}_{k}\equiv\bm{H}^{w}_{k}\bm{v}^{4}_{k}$ $\displaystyle\bm{H}^{w}_{k}\equiv\left[\bm{I}-\eta\gamma\bm{\Theta}\bm{C}_{3} \right]^{-1}\bm{\Theta}$ (27) $$

Now, combining this with the equation for $\bm{v}^{0}_{k}=\text{Vec}\{v^{0}_{k}(t)\}$ we find

$$ $\displaystyle\bm{v}^{0}_{k}=\bm{H}^{0}_{k}\left[w^{\star}_{k}\bm{1}-(\bm{H}^{w }_{k}+\eta\gamma\bm{C}_{w})\left(\bm{u}^{4}_{k}+\bm{R}_{3}\bm{u}^{2}_{k}\right )\right]$ $\displaystyle\bm{H}^{0}_{k}=\left[\bm{I}+\lambda_{k}\left(\bm{H}^{w}_{k}+\eta \gamma\bm{C}_{w}\right)\bm{R}_{3}\right]^{-1}$ (28) $$

The key response function is $\bm{R}_{3}$ with entries $[\bm{R}_{3}]_{t,s}=R_{3}(t,s)$ satisfies the following equation

$$ $\displaystyle\bm{R}_{3}=\bm{I}+\frac{1}{N}\bm{R}_{2,4}\bm{R}_{3}\ ,\ \bm{R}_{2 ,4}=-\sum_{k=1}^{M}\lambda_{k}\bm{H}^{0}_{k}\left(\bm{H}^{w}_{k}+\eta\gamma\bm {C}_{w}\right)$ $\displaystyle\implies$ $\displaystyle\bm{R}_{3}=\bm{I}-\frac{1}{N}\sum_{k=1}^{M}\lambda_{k}\bm{H}^{0}_ {k}\left(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w}\right)\bm{R}_{3}$ $\displaystyle=\bm{I}-\frac{1}{N}\sum_{k=1}^{M}\lambda_{k}\left[\bm{I}+\lambda_ {k}\left(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w}\right)\bm{R}_{3}\right]^{-1}\left (\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w}\right)\bm{R}_{3}$ (29) $$

Lastly, we can compute the correlation matrix $\bm{C}_{w}$ from the covariance $C_{3}$

$$ $\displaystyle\bm{C}_{w}=\left[\bm{I}-\eta\gamma\bm{C}_{3}\right]^{-1}\bm{ \Theta}\bm{C}_{3}\bm{\Theta}^{\top}\left[\bm{I}-\eta\gamma\bm{C}_{3}\right]^{- 1\top}$ (30) $$

The remaining correlation functions are defined as

$$ $\displaystyle\bm{C}_{0}$ $\displaystyle=\sum_{k}\lambda_{k}\bm{H}^{0}_{k}\left[(w^{\star}_{k})^{2}\bm{1} \bm{1}^{\top}+(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w})\left(\frac{1}{N}\bm{C}_{3} +\frac{\lambda_{k}}{B}\bm{R}_{3}\text{diag}(\bm{C}_{0})\bm{R}_{3}^{\top}\right )(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w})^{\top}\right][\bm{H}^{0}_{k}]^{\top}$ $\displaystyle\bm{C}_{2}$ $\displaystyle=\frac{1}{B}\sum_{k}\lambda_{k}\left(\bm{I}-\lambda_{k}\bm{H}^{0} _{k}(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w})\bm{R}_{3}\right)\text{diag}(\bm{C}_{ 0})\left(\bm{I}-\lambda_{k}\bm{H}^{0}_{k}(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w}) \bm{R}_{3}\right)^{\top}$ $\displaystyle+\sum_{k}\bm{H}_{k}^{0}\left[\bm{1}\bm{1}^{\top}(w_{k}^{\star})^{ 2}+\frac{1}{N}(\bm{H}^{w}_{k}+\eta\gamma\bm{C}_{w})\bm{C}_{3}(\bm{H}^{w}_{k}+ \eta\gamma\bm{C}_{w})^{\top}\right]\left[\bm{H}_{k}^{0}\right]^{\top}$ $\displaystyle\bm{C}_{3}$ $\displaystyle=\bm{R}_{3}\bm{C}_{2}\bm{R}_{3}^{\top}$ (31) $$

## Appendix D Offline Training: Train and Test Loss Under Sample Reuse

Our theory can also handle the case where samples are reused in a finite dataset $\bm{\Psi}\in\mathbb{R}^{P\times M}$. To simplify this setting we focus on the gradient flow limit (this will preserve all of the interesting finite-$P$ effects while simplifying the expressions)

$$ $\displaystyle\frac{d}{dt}\bm{w}(t)$ $\displaystyle=\bm{A}(t)\left(\frac{1}{P}\bm{\Psi}^{\top}\bm{\Psi}\right)\bm{v} ^{0}(t)$ $\displaystyle\frac{d}{dt}\bm{A}(t)$ $\displaystyle=\gamma\ \bm{w}(t)\bm{v}^{0}(t)^{\top}\left(\frac{1}{P}\bm{\Psi}^ {\top}\bm{\Psi}\right)\left(\frac{1}{N}\bm{A}(0)^{\top}\bm{A}(0)\right)$ (32) $$

We introduce the fields

$$ $\displaystyle\bm{v}^{1}(t)=\bm{\Psi}\bm{v}^{0}(t)\ ,\ \bm{v}^{2}(t)=\frac{1}{P }\bm{\Psi}^{\top}\bm{v}^{1}(t)$ (33) $\displaystyle\bm{v}^{3}(t)=\bm{A}(0)\bm{v}^{2}(t)\ ,\ \bm{v}^{4}(t)=\frac{1}{N }\bm{A}(0)^{\top}\bm{v}^{3}(t)$ (34) $$

so that the dynamics can be expressed as

$$ $\displaystyle\frac{d}{dt}\bm{w}(t)$ $\displaystyle=\bm{A}(t)\bm{v}^{2}(t)$ $\displaystyle\frac{d}{dt}\bm{A}(t)$ $\displaystyle=\gamma\bm{w}(t)\bm{v}^{4}(t)^{\top}$ (35) $$

As before we also introduce the following field which shows up in the $\bm{v}^{0}(t)$ dynamics

$$ $\displaystyle\bm{v}^{w}(t)=\frac{1}{N}\bm{A}(0)^{\top}\bm{w}(t)$ (36) $$

##### Data Average

The average over the frozen data matrix $\bm{\Psi}\in\mathbb{R}^{P\times M}$

$$ $\displaystyle v^{2}_{k}(t)=u^{2}_{k}(t)+\lambda_{k}\int_{0}^{t}dsR_{1}(t,s)v^{ 0}_{k}(s)\ ,\ u^{2}_{k}(t)\sim\mathcal{N}(0,P^{-1}\lambda_{k}C_{1}(t,s))$ (37) $\displaystyle v^{1}(t)=u^{1}(t)+\frac{1}{P}\int_{0}^{t}dsR_{0,2}(t,s)v^{1}(s) \ ,\ u^{1}(t)\sim\mathcal{N}(0,C_{0}(t,s))$ (38) $$

##### Feature Projection Average

Next we average over $\bm{A}(0)\in\mathbb{R}^{N\times M}$ with $N/M=\nu$ which yields

$$ $\displaystyle v^{4}_{k}(t)=u^{4}_{k}(t)+\int_{0}^{t}dsR_{3}(t,s)v^{2}_{k}(s)\ ,\ u^{4}_{k}(t)\sim\mathcal{N}(0,N^{-1}C_{3}(t,s))$ (39) $\displaystyle v^{3}(t)=u^{3}(t)+\frac{1}{N}\int_{0}^{t}dsR_{2,4}(t,s)v^{3}(s) \ ,\ u^{3}(t)\sim\mathcal{N}(0,C_{2}(t,s))$ (40) $$

We can now simply plug these equations into the dynamics of $\bm{w}(t),\bm{v}^{w}(t),\bm{v}^{0}(t)$ to obtain the final DMFT equations.

##### Final DMFT Equations for Data Reuse Setting

The complete governing equations for the test loss evolution after averaging over the random matrices $\{\bm{\Psi},\bm{A}\}$ can be obtained from the following stochastic processes which are driven by Gaussian noise sources $\{u^{2}_{k}(t),u^{3}(t),u^{4}_{k}(t)\}$. Letting $\left<\cdot\right>$ represent averages over these sources of noise, the equations close as

$$ $\displaystyle v^{0}_{k}(t)=w^{\star}_{k}-v^{w}_{k}(t)-\gamma\int_{0}^{t}dsC_{w }(t,s)v^{4}_{k}(s)$ $\displaystyle\partial_{t}v^{w}_{k}(t)=v^{4}_{k}(t)+\gamma\int_{0}^{t}dsC_{3}(t ,s)v^{w}_{k}(s)$ $\displaystyle v^{1}(t)=u^{1}(t)+\frac{1}{P}\int_{0}^{t}dsR_{0,2}(t,s)v^{1}(s) \ ,\ u^{1}(t)\sim\mathcal{N}(0,C_{0}(t,s))$ $\displaystyle v^{2}_{k}(t)=u^{2}_{k}(t)+\lambda_{k}\int dsR_{1}(t,s)v^{0}_{k}( t)\ ,\ u^{2}_{k}(t)\sim\mathcal{N}\left(0,P^{-1}\lambda_{k}C_{1}(t,s)\right)$ $\displaystyle v^{3}(t)=u^{3}(t)+\frac{1}{N}\sum_{s<t}R_{2,4}(t,s)v^{3}(s)\ ,\ u^{3}(t)\sim\mathcal{N}(0,C_{2}(t,s))$ $\displaystyle w(t+1)=w(t)+\eta v^{3}(t)+\eta^{2}\gamma\sum_{s<t}C_{3}(t,s)w(s)$ $\displaystyle v^{4}_{k}(t)=u^{4}_{k}(t)+\sum_{s<t}R_{3}(t,s)v^{2}_{k}(s)\ ,\ u ^{4}_{k}(t)\sim\mathcal{N}(0,N^{-1}C_{3}(t,s))$ $\displaystyle R_{0,2}(t,s)=\sum_{k}\lambda_{k}\left<\frac{\partial v^{0}_{k}(t )}{\partial u^{2}_{k}(s)}\right>\ ,\ R_{1}(t,s)=\left<\frac{\partial v^{1}(t)} {\partial u^{1}(s)}\right>$ $\displaystyle R_{2,4}(t,s)=\sum_{k}\left<\frac{\partial v^{2}_{k}(t)}{\partial u ^{4}_{k}(s)}\right>\ ,\ R_{3}(t,s)=\left<\frac{\partial v^{3}(t)}{\partial u^{ 3}(s)}\right>$ $\displaystyle C_{0}(t,s)=\sum_{k}\lambda_{k}\left<v^{0}_{k}(t)v^{0}_{k}(s) \right>\ ,\ C_{1}(t,s)=\left<v^{1}(t)v^{1}(s)\right>$ $\displaystyle C_{2}(t,s)=\sum_{k}\left<v^{2}_{k}(t)v^{2}_{k}(s)\right>\ ,\ C_{ 3}(t,s)=\left<v^{3}(t)v^{3}(s)\right>$ (41a) $$

The $\gamma\to 0$ limit of these equations recovers the DMFT equations from Bordelon et al. () which analyzed the random feature (static $\bm{A}$) case.

## Appendix E Bottleneck Scalings for Power Law Features

In this setting, we investigate the scaling behavior of the model under the source and capacity conditions described in the main text:

$$ $\displaystyle\lambda_{k}\sim k^{-\alpha}\ ,\ (w^{\star}_{k})^{2}\lambda_{k} \sim k^{-\beta\alpha-1}$ (42) $$

### E.1 Time Bottleneck

In this section we compute the loss dynamics in the limit of $N,B\to\infty$. We start with a perturbative argument that predicts a scaling law of the form $\mathcal{L}(t)\sim t^{-\beta(2-\beta)}$ for $\beta<1$. We then use this approximation to bootstrap more and more precise estimates of the exponent. The final prediction is the limit of infinitely many approximation steps which recovers $\mathcal{L}(t)\sim t^{-\frac{2\beta}{1+\beta}}$. We then provide a self-consistency derivation of this exponent to verify that it is the stable fixed point of the exponent.

#### E.1.1 Warmup: Perturbation Expansion of the DMFT Oder Parameters

First, in this section, we investigate the $N,B\to\infty$ limit of the DMFT equations and study what happens for small but finite $\gamma$. This perturbative approximation will lead to an approximation $\mathcal{L}\sim t^{-\beta(2-\beta)}$. In later sections, we will show how to refine this approximation to arrive at our self-consistently computed exponent $\frac{2\beta}{1+\beta}$. In the $N,B\to\infty$ limit, the DMFT equations simplify to

$$ $\displaystyle\bm{R}_{3}\to\bm{I}\ ,\ \bm{u}^{4}_{k}\to 0\ ,\ \bm{u}^{2}_{k}\to 0 \ ,\ \bm{v}^{4}_{k}\to\lambda_{k}\bm{v}^{0}_{k}\ ,\ \bm{C}_{3}\to\bm{C}_{2}.$ (43) $$

The dynamics in this limit have the form

$$ $\displaystyle v^{0}_{k}(t)=w^{\star}_{k}-v^{w}_{k}(t)-\eta\gamma\lambda_{k} \sum_{s<t}C_{w}(t,s)v^{0}_{k}(s)$ $\displaystyle v^{w}_{k}(t+1)-v^{w}_{k}(t)=\eta\lambda_{k}v^{0}_{k}(t)+\eta \gamma\sum_{s<t}C_{2}(t,s)v^{w}_{k}(s)$ $\displaystyle w(t+1)-w(t)=\eta v^{3}(t)+\eta\gamma\sum_{s<t}C_{2}(t,s)w(s)$ $\displaystyle C_{2}(t,s)=\sum_{k}\lambda_{k}^{2}\left<v^{0}_{k}(t)v^{0}_{k}(s) \right>\ ,\ C_{w}(t,s)=\left<w(t)w(s)\right>$ (44) $$

These exact dynamics can be simulated as we do in Figure 2.
However, we can obtain the correct rate of convergence by studying the following Markovian continuous time approximation of the above dynamics where we neglect the extra $\mathcal{O}(\gamma)$ term in the $v^{w}_{k}(t)$ dynamics

$$ $\displaystyle\frac{d}{dt}v^{0}_{k}(t)\approx\lambda_{k}v^{0}_{k}(t)-\gamma \lambda_{k}C_{w}(t)v^{0}_{k}(t)\ ,\ C_{w}(t)\equiv C_{w}(t,t)$ $\displaystyle\partial_{t}C_{w}(t)\approx C_{2}(t)+\mathcal{O}(\gamma)\ ,\ C_{2 }(t)\equiv C_{2}(t,t)$ (45) $$

The solution for the error along the $k$-th eigenfunction will take the form

$$ $\displaystyle v^{0}_{k}(t)\approx\exp\left(-\lambda_{k}t-\gamma\lambda_{k}\int _{0}^{t}dsC_{w}(s)\right)w^{\star}_{k}$ (46) $$

We next solve for the dynamics of $C_{2}(t)$ and $C_{w}(t)$ in the leading order $\gamma\to 0$ limit (under the linear dynamics)

$$ $\displaystyle C_{2}(t)\sim\sum_{k}\lambda_{k}^{2}(w^{\star}_{k})^{2}e^{-2 \lambda_{k}t}\sim\int dkk^{-\alpha-\beta\alpha-1}e^{-k^{-\alpha}t}\sim t^{-\beta}$ $\displaystyle C_{w}(t)\sim\begin{cases}t^{1-\beta}&\beta<1\\ C_{w}(\infty)&\beta>1\end{cases}$ (47) $$

where $C_{w}(\infty)$ is a limiting finite value of $C_{w}(t)$.

$$ $\displaystyle\frac{v^{0}_{k}(t)}{w^{\star}_{k}}\approx\begin{cases}\exp\left(- \lambda_{k}t-\gamma\lambda_{k}t^{2-\beta}\right)&\beta<1\\ \exp\left(-\lambda_{k}[1+\gamma C_{w}(\infty)]t\right)&\beta>1\end{cases}$ (48) $$

For $\beta<1$, the feature learning term will eventually dominate. The mode $k_{\star}(t)$ which is being learned at time $t$ satisfies

$$ $\displaystyle k_{\star}(t)\sim t^{(2-\beta)/\alpha}$ (49) $$

which implies that the

$$ $\displaystyle\mathcal{L}\approx\sum_{k>k_{\star}}(w^{\star}_{k})^{2}\lambda_{k }\sim t^{-\beta(2-\beta)}$ (50) $$

However, this solution actually over-estimates the exponent. To derive a better approximation of the exponent, we turn to a Markovian perspective on the dynamics which holds as $N,B\to\infty$.

#### E.1.2 Matrix Perturbation Perspective on the Time Bottleneck Scaling with Markovian Dynamics

The limiting dynamics in the $N,B\to\infty$ limit can also be expressed as a Markovian system in terms of the vector $\bm{v}^{0}(t)$ and a matrix $\bm{M}(t)$ which preconditions the gradient flow dynamics

$$ $\displaystyle\frac{d}{dt}\bm{v}^{0}(t)=-\bm{M}(t)\bm{\Lambda}\bm{v}^{0}(t)\ , \ \bm{M}(t)=\left[\frac{1}{N}\bm{A}(t)^{\top}\bm{A}(t)+\frac{\gamma}{N}|\bm{w} (t)|^{2}\bm{I}\right]$ $\displaystyle\frac{d}{dt}\bm{M}(t)=\gamma\left(\bm{w}_{\star}-\bm{v}^{0}(t) \right)\bm{v}^{0}(t)^{\top}\bm{\Lambda}+\gamma\bm{\Lambda}\bm{v}^{0}(t)\left( \bm{w}_{\star}-\bm{v}^{0}(t)\right)^{\top}+2\gamma(\bm{w}_{\star}-\bm{v}^{0}(t ))^{\top}\bm{\Lambda}\bm{v}^{0}(t)\bm{I}$ (51) $$

We can rewrite this system in terms of the function $\bm{\Delta}(t)=\bm{\Lambda}^{1/2}\bm{v}^{0}(t)$ with $\bm{y}=\bm{\Lambda}^{1/2}\bm{w}_{\star}$ and the Hermitian kernel matrix $\bm{K}=\bm{\Lambda}^{1/2}\bm{M}(t)\bm{\Lambda}^{1/2}$ then

$$ $\displaystyle\frac{d}{dt}\bm{\Delta}(t)=-\bm{K}(t)\bm{\Delta}(t).$ $\displaystyle\frac{d}{dt}\bm{K}(t)=\gamma(\bm{y}-\bm{\Delta}(t))\bm{\Delta}(t) ^{\top}\bm{\Lambda}+\gamma\bm{\Lambda}\bm{\Delta}(t)(\bm{y}-\bm{\Delta}(t))^{ \top}+2\gamma(\bm{y}-\bm{\Delta}(t))\cdot\bm{\Delta}(t)\ \bm{\Lambda}$ (52) $$

The test loss can be expressed as $\mathcal{L}(t)=|\bm{\Delta}(t)|^{2}$.

##### Loss Dynamics Dominated by the Last Term in Kernel Dynamics

We note that the loss dynamics satisfy the following dynamics at large time $t$

$$ $\displaystyle\frac{d}{dt}\mathcal{L}(t)=-\bm{\Delta}(t)^{\top}\bm{K}(t)\bm{ \Delta}(t)$ $\displaystyle=-\bm{\Delta}(t)\bm{\Lambda}\bm{\Delta}(t)-2\gamma\int_{0}^{t}ds[ (\bm{y}-\bm{\Delta}(s))\cdot\bm{\Delta}(t)]\bm{\Delta}(t)^{\top}\bm{\Lambda} \bm{\Delta}(s)$ $\displaystyle-2\gamma(\bm{\Delta}(t)^{\top}\bm{\Lambda}\bm{\Delta}(t))\int_{0} ^{t}ds(\bm{y}-\bm{\Delta}(s))\cdot\bm{\Delta}(s)$ $\displaystyle\sim-\underbrace{\bm{\Delta}(t)^{\top}\bm{\Lambda}\bm{\Delta}(t)} _{\text{Lazy Limit}}-\underbrace{2\gamma[\bm{y}\cdot\bm{\Delta}(t)]\int_{0}^{t }ds\bm{\Delta}(t)^{\top}\bm{\Lambda}\bm{\Delta}(s)}_{\text{Subleading}}- \underbrace{2\gamma(\bm{\Delta}(t)^{\top}\bm{\Lambda}\bm{\Delta}(t))\int_{0}^{ t}ds\ \bm{y}\cdot\bm{\Delta}(s)}_{\text{Dominant}}$ (53) $\displaystyle\approx-\bm{\Delta}(t)^{\top}\bm{\Lambda}\bm{\Delta}(t)\left[1+2 \gamma\int_{0}^{t}ds\bm{y}\cdot\bm{\Delta}(s)\right].$ (54) $$

One can straightforwardly verify that the middle term is subleading compared to the final term is that under the ansatz that $\Delta_{k}(t)\sim\exp\left(-\lambda_{k}t^{2-\chi}\right)$ where $\beta\leq\chi\leq 1$ for $\beta<1$. We can therefore focus on the last term when deriving corrections to the scaling law.

##### Intuition Pump: Perturbative Level 1 Approximation

In the lazy limit $\gamma\to 0$, $\bm{K}(t)=\bm{\Lambda}$ for all $t$. However, for $\gamma>0$ this effective kernel matrix $\bm{K}(t)$ evolves in a task-dependent manner. To compute $\bm{K}$ will approximate $\bm{M}(t)$ with its leading order dynamics in $\gamma$, which are obtained by evaluating the $\bm{v}^{0}(t)$ dynamics with the lazy learning $\gamma\to 0$ solution. We can thus approximate the kernel matrix $\bm{K}(t)$ dynamics as

$$ $\displaystyle\bm{K}(t)$ $\displaystyle\approx\bm{\Lambda}+\gamma\bm{y}\bm{y}^{\top}\left(\bm{I}-\exp \left(-\bm{\Lambda}t\right)\right)^{\top}+\gamma\left(\bm{I}-\exp(-\bm{\Lambda }t)\right)\bm{y}\bm{y}^{\top}$ $\displaystyle+2\gamma\left[\bm{y}^{\top}\bm{\Lambda}^{-1}\left(\bm{I}-\exp(- \bm{\Lambda}t)\right)\bm{y}\right]\bm{\Lambda}$ (55) $$

From this perspective we see that the kernel has two dynamical components. First, a low rank spike grows in the kernel, eventually converging to the rank one matrix $\bm{y}\bm{y}^{\top}$. In addition, there is a scale growth of the existing eigenvalues due to the last term $\left[\bm{y}^{\top}\bm{\Lambda}^{-1}\left(\bm{I}-\exp(-\bm{\Lambda}t)\right)
\bm{y}\right]\bm{\Lambda}$, which will approach the value of the RKHS norm of the target function as $t\to\infty$. The eigenvalues $\{\mathcal{K}_{k}(t)\}_{k=1}^{\infty}$ of the kernel $\bm{K}(t)$ evolve at leading order as the diagonal entries. Assuming that $\beta<1$ these terms increase with $t$ as

$$ $\displaystyle\mathcal{K}_{k}(t)$ $\displaystyle\sim\lambda_{k}+2\gamma y_{k}^{2}\left(1-e^{-\lambda_{k}t}\right) +2\gamma\lambda_{k}\sum_{\ell}\frac{y_{\ell}^{2}}{\lambda_{\ell}}(1-e^{- \lambda_{\ell}t})$ $\displaystyle\sim\lambda_{k}+2\gamma y_{k}^{2}\left(1-e^{-\lambda_{k}t}\right) +2\gamma\lambda_{k}t^{1-\beta}$ (56) $$

The dynamics for the errors can be approximated as

$$ $\displaystyle\frac{\partial}{\partial t}\Delta_{k}(t)\sim-\mathcal{K}_{k}(t) \Delta_{k}(t)$ $\displaystyle\implies\Delta_{k}(t)\sim\exp\left(-\int_{0}^{t}ds\ \mathcal{K}_{ k}(s)\right)\sqrt{\lambda_{k}}w^{\star}_{k}$ $\displaystyle\sim\exp\left(-\lambda_{k}t-2\gamma\lambda_{k}(w^{\star}_{k})^{2} t-2\gamma\lambda_{k}t^{2-\beta}\right)\sqrt{\lambda_{k}}w^{\star}_{k}$ (57) $$

For sufficiently large $t$, the final term dominates and the mode $k_{\star}(t)$ which is being learned at time $t$ is

$$ $\displaystyle k_{\star}(t)\sim t^{\frac{2-\beta}{\alpha}}$ (58) $$

The test loss is simply the variance in the unlearned modes

$$ $\displaystyle\mathcal{L}(t)\sim\sum_{k>k_{\star}}(w^{\star}_{k})^{2}\lambda_{k }\sim t^{-\beta(2-\beta)}.$ (59) $$

##### Bootstrapping a More Accurate Exponent

From the previous argument, we started with the lazy learning limiting dynamics for $v^{0}_{k}(t)\sim e^{-\lambda_{k}t}w^{\star}_{k}$ and used these dynamics to estimate
the rate at which $M(t)$ (or equivalently $K(t)$) changes. This lead to an improved rate of convergence for the mode errors $v^{0}_{k}(t)$, which under this next order approximation decay as $v^{0}_{k}(t)\sim e^{-\lambda_{k}t^{2-\beta}}w^{\star}_{k}$. Supposing that the errors decay at this rate, we can estimate the dynamics of $M(t)$

$$ $\displaystyle\frac{d}{dt}\mathcal{K}_{k}(t)\approx\lambda_{k}\sum_{\ell}(w^{ \star}_{\ell})^{2}\lambda_{\ell}e^{-\lambda_{\ell}t^{2-\beta}}\approx\lambda_{ k}t^{-\beta(2-\beta)}\ ,\$ (60) $\displaystyle\implies v^{0}_{k}(t)\sim\exp\left(-\lambda_{k}t^{2-\beta(2-\beta )}\right)w^{\star}_{k}\ ,\ (\text{Level 2 Approximation})$ (61) $$

We can imagine continuing this approximation scheme to higher and higher levels which will yield a series of better approximations to the power law

$$ $\displaystyle\mathcal{L}(t)\sim\begin{cases}t^{-\beta}&\text{Level 0 Approximation}\\ t^{-\beta(2-\beta)}&\text{Level 1 Approximation}\\ t^{-\beta\left[2-\beta(2-\beta)\right]}&\text{Level 2 Approximation}\\ t^{-\beta\left[2-\beta(2-\beta(2-\beta))\right]}&\text{Level 3 Approximation} \\ ...\end{cases}$ (62) $$

We plot the first few of these in Figure 10, showing that they approach a limit as the number of levels diverges. As $n\to\infty$, this geometric series will eventually converge to $t^{-\frac{2\beta}{1+\beta}}$.

Figure: Figure 10: Predictions for the loss scaling exponent $\chi_{n}(\beta)$ at varying levels $n$ of the approximation scheme. Our final prediction is the infinite level limit which gives $\frac{2}{1+\beta}$. This agrees with a self-consistency argument.
Refer to caption: x23.png

### E.2 Self-Consistent Derivation of the Infinite Level Scaling Law Exponent

From the above argument, it makes sense to wonder whether or not there exists a fixed point to this series of approximations that will actually yield the correct exponent in the limit of infinitely many steps. Indeed in this section, we find that this limit can be computed self-consistently

$$ $\displaystyle\frac{d}{dt}M(t)\approx\gamma\lambda_{k}\sum_{\ell}w^{\star}_{ \ell}\lambda_{\ell}v^{0}_{\ell}(t)$ (63) $\displaystyle\implies M(t)=1+\gamma\int_{0}^{t}ds\sum_{\ell}w^{\star}_{\ell} \lambda_{\ell}v_{\ell}^{0}(t)$ (64) $\displaystyle v_{k}(t)\sim\exp\left(-\lambda_{k}\int_{0}^{t}dt^{\prime}M(t^{ \prime})\right)w^{\star}_{k}$ (65) $$

We can define the intermediate variable

$$ $\displaystyle B(t)=\sum_{k}w^{\star}_{k}\lambda_{k}v_{k}^{0}(t)$ (66) $$

which has self-consistent equation

$$ $\displaystyle B(t)=\sum_{k}\lambda_{k}(w^{\star}_{k})^{2}\exp\left(-\lambda_{k }\int_{0}^{t}dt^{\prime}\left[1+\gamma\int_{0}^{t^{\prime}}dsB(s)\right]\right)$ (67) $$

We can seek a solution of the form $B(t)\sim t^{-\chi}$. This yields

$$ $\displaystyle t^{-\chi}\approx t^{-\max\{\beta,\beta(2-\chi)\}}\implies\chi= \beta\max\left\{1,\frac{2}{\beta+1}\right\}.$ (68) $$

Using the solution for $B(t)$, we can also derive the scaling for the loss which is identical $\mathcal{L}(t)\sim t^{-\chi}$. We note that this argument also leads to an approximate doubling of the exponent compared to the lazy case for $\beta<1$, however this is slightly disagreement with the perturbative approach which yields $\beta(2-\beta)$ for $\beta<1$.

This argument is where we developed the general expression that determines $\chi$ which was provided in the main text

$$ $\displaystyle\chi=-\lim_{t\to\infty}\frac{1}{\ln t}\ln\left[\sum_{k}(w^{\star} _{k})^{2}\lambda_{k}\exp\left(-\lambda_{k}\left[t+\gamma t^{2-\chi}\right] \right)\right].$ (69) $$

### E.3 Finite Model Size Bottleneck

In the limit of $t\to\infty$, the dynamics for $\{v^{0}_{k}(t)\}$ will converge to a fixed point that depends on $N$. To ascertain the value of this fixed point, we first must compute the asymptotics. First, we note that correlation and response functions reach the following fixed point behaviors

$$ $\displaystyle\lim_{t,s\to\infty}\int_{0}^{t}dt^{\prime}R_{3}(t^{\prime},s)\sim \ r_{3}\delta(t-s)$ $\displaystyle\lim_{t,s\to\infty}R_{2,4}(t,s)=r_{2,4}\ \Theta(t-s)$ $\displaystyle\lim_{t,s\to\infty}C_{w}(t,s)=c_{w}$ $\displaystyle\lim_{t,s\to\infty}\int_{0}^{t}dt^{\prime}C_{3}(t^{\prime},s)=0$ (70) $$

which gives the following long time behavior

$$ $\displaystyle v^{0}_{k}(t)$ $\displaystyle\sim w^{\star}_{k}-\int_{0}^{t}dt^{\prime}u^{4}_{k}(t^{\prime})- \lambda_{k}\int_{0}^{t}dt^{\prime}\int_{0}^{t^{\prime}}dsR_{3}(t^{\prime},s)v^ {0}_{k}(s)$ (71) $\displaystyle\sim w^{\star}_{k}-\int_{0}^{t}dt^{\prime}u^{4}_{k}(t^{\prime})- \lambda_{k}r_{3}v^{0}_{k}(t)$ (72) $\displaystyle\implies r_{2,4}$ $\displaystyle\sim-\sum_{k}\frac{\lambda_{k}}{1+\lambda_{k}r_{3}}$ (73) $$

Using the asymptotic relationship between $\frac{1}{N}r_{3}r_{2,4}=\nu$, we arrive at the following self-consistent equation for $r_{3}$

$$ $\displaystyle 1=\frac{1}{N}\sum_{k}\frac{\lambda_{k}r_{3}}{1+\lambda_{k}r_{3}}$ (74) $$

For power law features this gives

$$ $\displaystyle N\approx\int dk\frac{k^{-\alpha}r_{3}}{k^{-\alpha}r_{3}+1} \approx[r_{3}]^{1/\alpha}\implies r_{3}\sim N^{\alpha}$ (75) $$

which recovers the correct scaling law with model size $N$.

$$ $\displaystyle\lim_{t,s\to\infty}C_{0}(t,s)$ $\displaystyle=\sum_{k}\frac{\lambda_{k}(w^{\star}_{k})^{2}}{(1+\lambda_{k}N^{ \alpha})^{2}}$ $\displaystyle\approx N^{-2\alpha}\int_{1}^{N}dk\ k^{-\alpha\beta-1+2\alpha}+ \int_{N}^{\infty}dkk^{-\beta\alpha-1}\sim N^{-\alpha\min\{2,\beta\}}$ (76a) $$

### E.4 Finite Data Bottleneck (Data Reuse Setting)

The data bottleneck when training on repeated $P$ training examples is very similar. In this case, the relevant response functions to track are $R_{1}(t,s)$ and $R_{0,2}(t,s)$ which have the following large time properties as $t,s\to\infty$

$$ $\displaystyle\int_{0}^{t}dt^{\prime}R_{1}(t^{\prime},s)\sim\ r_{1}\delta(t-s)$ $\displaystyle R_{0,2}(t,s)\sim r_{0,2}\ \Theta(t-s)$ $$

Under this ansatz, we find the following expression for $r_{1}$ as $t\to\infty$

$$ $\displaystyle 1\sim\frac{1}{P}\sum_{k}\frac{\lambda_{k}r_{1}}{1+\lambda_{k}r_{ 1}}.$ (77) $$

Following an identical argument above we find that $r_{1}\sim P^{\alpha}$, resulting in the following asymptotic test loss

$$ $\displaystyle\lim_{t,s\to\infty}C_{0}(t,s)$ $\displaystyle=\sum_{k}\frac{\lambda_{k}(w^{\star}_{k})^{2}}{(1+\lambda_{k}P^{ \alpha})^{2}}$ $\displaystyle\approx P^{-2\alpha}\int_{1}^{P}dkk^{-\alpha\beta-1+2\alpha}+\int _{P}^{\infty}dkk^{-\beta\alpha-1}\sim P^{-\alpha\min\{2,\beta\}}.$ (78a) $$

## Appendix F Transient Dynamics

To compute the transient $1/N$ and $1/B$ effects, it suffices to compute the scaling of a response function / Volterra kernel at leading order.

### F.1 Leading Bias Correction at Finite N 𝑁 N italic_N

At leading order in $1/N$ the bias corrections from finite model size can be obtained from the following leading order approximation of the response function $R_{3}(t,s)$

$$ $\displaystyle R_{3}(t,s)$ $\displaystyle\sim\delta(t-s)-\frac{1}{N}\sum_{k}\lambda_{k}e^{-\lambda_{k}(t^{ \chi/\beta}-s^{\chi/\beta})}+\mathcal{O}(N^{-2})$ $\displaystyle\sim\delta(t-s)+\frac{1}{N}(t^{\chi/\beta}-s^{\chi/\beta})^{-1+1/\alpha}$ (79) $$

where $\chi=\beta\max\left\{1,\frac{2}{1+\beta}\right\}$. Following Paquette et al. (), we note that the scaling of $R_{3}(t,s)-\delta(t-s)$ determines the scaling of the finite width transient. Thus the finite width effects can be approximated as

$$ $\displaystyle\mathcal{L}(t,N)\sim\underbrace{t^{-\beta\max\{1,\frac{2}{1+\beta }\}}}_{\text{Limiting Dynamics}}+\underbrace{N^{-\alpha\min\{2,\beta\}}}_{ \text{Model Bottleneck}}+\underbrace{\frac{1}{N}t^{-(1-1/\alpha)\max\{1,\frac{ 2}{1+\beta}\}}}_{\text{Finite Model Transient}}.$ (80) $$

In the case where $\beta>1$ this agrees with the transient derived in Paquette et al. (). However for $\beta<1$, the feature learning dynamics accelerate the decay rate of this term.

### F.2 Leading SGD Correction at Finite B 𝐵 B italic_B

We start by computing the $N\to\infty$ limit of the SGD dynamics can be approximated as

$$ $\displaystyle C_{0}(t)$ $\displaystyle\approx\sum_{k}(w^{\star}_{k})^{2}\lambda_{k}\exp\left(-2\lambda_ {k}(t+\gamma t^{2+\chi})\right)+\frac{\eta}{B}\int_{0}^{t}dsK(t,s)C_{0}(s)$ (81) $$

where $\chi=\beta\max\left\{1,\frac{2}{1+\beta}\right\}$ and $K(t,s)$ is the Volterra-kernel for SGD Paquette et al. (; ), which in our case takes the form

$$ $\displaystyle K(t,s)$ $\displaystyle=\sum_{k}\lambda_{k}\exp\left(-2\lambda_{k}\left[(t+\gamma t^{2- \chi})-(s-\gamma s^{2-\chi})\right]\right)$ $\displaystyle\sim\left(t^{\frac{\max(1,2-\chi)}{\alpha}}-s^{\frac{\max(1,2- \chi)}{\alpha}}\right)^{-(\alpha-1)}$ (82) $$

Since the transient dynamics are again generated by the long time behavior of $K(t,0)$, we can approximate the SGD dynamics as

$$ $\displaystyle\mathcal{L}(t,B)\approx\underbrace{t^{-\beta\max\{1,\frac{2}{1+ \beta}\}}}_{\text{Gradient Flow}}+\underbrace{\frac{\eta}{B}t^{-(1-1/\alpha) \max\left\{1,\frac{2}{1+\beta}\right\}}}_{\text{SGD Noise}}.$ (83) $$

As before, the $\beta>1$ case is consistent with the estimate for the Volterra kernel scaling in Paquette et al. ().

## Appendix G Linear Network Dynamics under Source and Capacity

In this section, we show how in a simple linear network, the advantage in the scaling properties of the loss due to larger $\gamma$ is evident. Here, we consider a simple model of a two-layer linear neural network trained with vanilla SGD online. We explicitly add a feature learning parameter $\gamma$. We study when this linear network can outperform the rate of $t^{-\beta}$ given by linear regression directly from input space. Despite its linearity, this setting is already rich enough to capture many of the power laws behaviors observed in realistic models.

### G.1 Model Definition

Following Chizat et al. (), the network function is parameterized as:

$$ $f(\bm{x};t)=\frac{1}{\gamma}(\tilde{f}(\bm{x};t)-\tilde{f}(\bm{x};0)),\quad \tilde{f}(\bm{x};t)=\bm{w}^{\top}\bm{A}\bm{x},\quad$ (84) $$

We let $\bm{x}\in\mathbb{R}^{M}$ and take hidden layer width $N$ so that $\bm{A}\in\mathbb{R}^{N\times M}$, $\bm{w}\in\mathbb{R}^{N}$, as in the main text.

We train the network on power law data of the following form

$$ $\bm{x}\sim\mathcal{N}(0,\bm{\Lambda}),\quad y=\bm{w}^{*}\cdot\bm{x}.$ (85) $$

We impose the usual source and capacity conditions on $\bm{\Lambda}$ and $\bm{w}^{*}$ as in equation 6.

### G.2 Improved Scalings only Below β < 1 𝛽 1 \beta<1 italic_β < 1

We empirically find that when $\beta>1$, large $\gamma$ networks do not achieve better scaling than small $\gamma$ ones. By contrast, when $\beta<1$ we see an improvement to the loss scaling. We illustrate both of these behaviors in Figure 11. Empirically, we observe a rate of $t^{-2\beta/(1+\beta)}$ for these linear networks. This is the same improvement derived for the projected gradient descent model studied in the main text.

Figure: Figure 11: Linear Networks a) $\beta>1$, where across values of $\gamma$, we observe the same asymptotic scaling going as $t^{-\beta}$ as predicted by kernel theory. b) $\beta<1$, where feature learning linear networks achieve an improved scaling as predicted by our theory.
Refer to caption: x24.png

### G.3 Tracking The Rank One Spike

In the linear network setting, one can show that this improved scaling is due to the continued the growth of a rank one spike in the first layer weights $\bm{A}$ of the linear network. By balancing, as in Du et al. (), this is matched by the growth of $|\bm{w}|^{2}$. This which will continue to grow extensively in time $D$ only when $\beta<1$. We illustrate these two cases in Figure 12.

Figure: Figure 12: We study the growth of the spike as measured by $|\bm{w}|^{2}$. (a) For easy tasks, where $\bm{w}^{*}$ is finite, the spike grows to finite size, and then plateaus. This leads to a multiplicative improvement in the loss, but does not change the scaling exponent. (b) When the task is hard, $\bm{w}$ continues to grow without bound. Both the perturbative scaling of $t^{1-\beta}$ and the scaling $t^{1-2\beta(1+\beta)}$ obtained from the self-consistent equation 10 are plotted. We see excellent agreement with the latter scaling.
Refer to caption: x26.png

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: The effect observed in this paper differs significantly from other feature learning paper in that it does not assume any form of sparse structure (actually feature learning appears to help when the signal is closest to being of the same intensity along all directions), i.e. low-index models for shallow networks, or low-rank matrices for linear networks. This suggest that the type of feature learning observed here is quite different from these other line of results. My intuition from your paper, is that feature learning may in some sense improve the conditioning of the Hessian/NTK, allowing faster convergence, whereas in the presence of sparsity feature learning seems to rather make the NTK low-rank and more ill-conditioned. I was therefore wondering whether the effects observed here might disappear when switching from online learning to GD on a fixed dataset, because it would untangle the training speed from the number of datapoints observed, and maybe we could see there that feature learning makes convergence faster but maybe one still converges to approximately the same solution (and thus same test error).

Typos:
- on page 5: The sentence "The first terms represent bottleneck/resolution-limited scalings in the sense that taking all other quantities to infinity and studying scaling with the last ..." does not finish.
soundness: 4
strengths: The paper uses advanced theoretical tools to tackle the important question of the effect of feature learning. A complete phase diagram is obtained, showing under which condition feature learning leads to a significant advantage. The theoretical scaling laws are checked empirically, and seem to even extend to some degree to deep nonlinear networks (which are not covered by the theory).
summary: The paper does a mean-field analysis of feature learning in a model similar to a shallow linear network. Scaling laws are obtained, where depending on the source exponent, feature learning can lead to a improvement in the exponent.
weaknesses: The paper does not give much intuition for why this improvement of feature learning is observed, instead the authors simply mention the existence of a self-consistent dynamic derived in the appendix.

I also feel that the authors do not relate much to the already existing litterature on the dynamics of linear networks (e.g. https://arxiv.org/abs/2211.16980), where feature learning and self-consistent dynamics have been described in https://arxiv.org/abs/2406.06158 or https://arxiv.org/abs/2405.17580 . It would be interesting to know whether the dynamics observed in this paper is similar to the ones observed in these papers.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: See weaknesses
soundness: 2
strengths: 1. Tractable Theoretical Setting: The paper presents a solvable theoretical model, which allows for potentially clearer and rigorous exploration of neural scaling laws beyond the kernel regime. This tractable setting potentially makes it easier to understand and predict the dynamics of feature learning in neural networks.

2. Extension of Prior Work: The authors improve upon simpler one-layer linear models from prior work by introducing a more sophisticated two-layer model. This extension potentially can provides deeper insights into the scaling behavior of neural networks.
summary: The paper introduces a solvable two-layer linear model trained using a variant of stochastic gradient descent (SGD) to predict the scaling laws beyond the kernel (lazy training) regime. The authors identify three scaling regimes—hard, easy, and super easy tasks—under which feature learning can improve the scaling of the loss with respect to training time and compute. The paper also provides experimental validation of these theoretical predictions by training deep, nonlinear models to demonstrate the applicability of the scaling laws.
weaknesses: 1. Complexity and Assumptions: The paper is difficult to follow due to the heavy reliance on prior work. This reliance can make it challenging for readers who are not deeply familiar with the specific foundational literature to fully grasp the contributions.

2. Unclear Definition of Feature Learning: Throughout the paper, the concept of "feature learning" is not clearly defined. It remains unclear how the authors measure feature learning or provide a tangible understanding of it. Although allowing 
A(t) to evolve leads to better performance, this change doesn’t adequately explain how features are being learned or evolved.

3. Limited Model Scope: The model seems quite limited in its design. Allowing 
A(t) to change is reasonable, especially since the target function is 𝑊∗𝜓_{inf}(x). However, it is unclear why this is believed to represent or mimic feature learning in practical, real-world scenarios. The connection to real feature learning remains vague. This is also related to 2nd point, that I do not understand how authors define feature learning.

4. Transferability to Deep Nonlinear Models: It is not clear why the results derived from the proposed linear model are expected to transfer well to deep, nonlinear models. 

overall I feel like the paper is scattered and I can not grasp the main point authors are trying to make about feature learning.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: From the experimental results in Figures 9 and 12, we can see that the scaling law does not change suddenly at $\gamma = 0$, but changes gradually at a very small value of gamma. Can this paper's analysis cover this aspect?
soundness: 3
strengths: 1. This paper constructs a comprehensive scaling law framework, reveals the power law relationship between loss and the number of parameters, training time, and batch size, and provides a theoretical basis for understanding the performance of neural networks.

2. This paper conducts a series of experiments to verify the theoretical results, making the conclusions of this paper more convincing.
summary: This paper studies the scaling laws of neural networks in two regimes: lazy learning and feature learning. Eq (8) and (9) describe the power-law variation of the loss with training time \( t \), number of parameters \( N \), and batch size \( B \) in these two regimes which show that feature learning accelerates the convergence of hard tasks.
weaknesses: 1. This paper ignores the impact of gamma on the feature learning scaling law.

2. The paper lacks an analysis of the critical state between the two regimes.

3. (Minor) The colors of lines in the figures are not differentiated enough.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: In line 160 the "allow" pops a bit out of nowhere and it hard to grasp. Following with Lines 175-176, my understanding is that you choose an arbitrary kind of feature learning dynamics, with the implicit assumption that SDG is sufficient (but not necessary) to obtain such dynamics. Being sufficient means that the scaling law you derive is at least an upper bound on the loss. While not being necessary means that a different choice of modelling dynamics for $A(t)$ may lead to tighter scaling low (which may for example explain the CIFAR experiment). Is this understanding correct? Is such implicit assumption true? Can you please argument a bit more over this choice of feature learning dynamics?

The reference "Bordelon et al. (2024a)" appears very often in the paper. It would be nice to have a short paragraph (maybe in the appendix) to clarify the common point with such paper and the novel contribution of this one over that one. Can you reply with a draft of such paragraph?
soundness: 4
strengths: The paper is very well written and easy to follow, overall very pleasant to read.
Settings and limitations are well clarified: the authors are very open about the limits of their theory and the non-success of the CIFAR experiment. This behaviour is super valuable and very helpful for future works and possible improvements.
Derivations are explained both intuitively and formally: key concepts are present in the main paper with a good trade-off between formalism and explanation; technical details not directly useful to understand the concepts are properly deferred to the appendix.
Results are clear and well discussed: for example the scaling laws are divided in several components, each of them is further discussed and empirically studied with ad-hoc experiments.
summary: The paper studies scaling laws of the training loss with respect to relevant parameters (training time, model size and batch size), with a focus on online training, where dataset size is assumed to be unlimited.
They study training dynamics through the lens of the eigenspace decomposition of the NTK. Such eigenspace is used both to (1) formally define capacity and source of model and task and (2) to express the loss behaviour explicitly, both in the lazy learning regime ($\gamma=0$) and in the feature learning regime ($\gamma>0$), the former can be solved explicitly while the second is approximated with DMFT. Such analysis reveals two very different regimes depending on the value of the source, specifically in the hard regime (source<1) they prove better scaling laws than state of the art. Such scaling laws are studied and discussed in each component. Empirically are shown to be tight on a variety of task, including synthetic and MNIST vision, but they are shown to no be tight in the CIFAR case.
weaknesses: The paper is overall solid and I don't see major weakness.

A big assumption made is the online setting, and the extension to finite dataset setting in Appendix D is somehow limited. However such kind of limitation has to be expected from a theoretical paper. Of course a general case analysis would be preferable, but that would very ambitious and I don't think this assumption can be considered a major weakness.

The dynamics mean field theory is not at all explained in the main paper, understandably. An informal explanation attempt is done in Line 202-211. While I consider this attempt valuable, the Equation (in line 206) is too much out of context, the notation is not defined and it ends up being practically useless for the reader. I recommend either removing (and using that space to expand the intuition informal explanation) or extending it, explaining the terms involved.

Minor detail. In Line 133 "we then diagonalize..." can be misinterpreted as "we discard the non-diagonal elements". I think overall the message is clear because right after (in Line 134) you refer to eigenfunctions, but I'd recommend trying to reformulate this sentence to avoid ambiguities.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Partially critical questions ordered from most to least important:

1. **Large feature learning regime.** Which $\gamma_0$ is optimal and is this another hyperparameter that has to be tuned, or are there useful heuristics for choosing it? In your experiments it looks like the larger the better, but the $\gamma_0$ you consider are still pretty small. I am guessing in Figure 6 (b), larger values of $\gamma_0$ would result in a better than predicted scaling exponent even earlier in training? If so, this would be an acceptable but important limitation of the studied model, which should be acknowledged. I suggest running the experiment for larger $\gamma_0$ and discussing this limitation more deeply.
2. **Importance of constants.** While reporting the pure scalings is nice for readability, in practice, how important/huge are the omitted constants? This could arbitrarily change the global optimum of the test error given finite resources. Towards answering this question of how practically useful the results are, I suggest an experiment that compares the predicted compute-optimal scaling to empirical loss measurements like in Figure 4 but for the more practical datasets in Figure 6.
3. **Error of the approximation.** How good is your approximation in eq. (9)? In which cases does it fail?
4. **Extension to label noise.** How would label noise affect the results?
5. **Long training regime.** What is happening under long training in Figure 6b? Do you have an informed guess? Are structures in the images learned that are not covered by your simple distributional assumption?
6. In eq. (6), you state $\lambda_k\sim k^{-\alpha}$ and $\lambda_k^{-\beta}\sim k^{-\alpha \beta}$. Is this a typo?
7. In l. 192, ‘the components of $\omega^*, v^0(t)$ in the k-th eigenspace of $\Lambda$’, but $\Lambda$ was defined as a diagonal matrix in eq. (2). I guess here the k-th eigenspace of $K_\infty$ is meant?
soundness: 2
strengths: The paper is very well written and easy to read. To the best of my knowledge, this is the first paper that theoretically explains compute-optimal scaling laws in non-linear feature learning networks beyond the lazy regime. Separating the four bottleneck mechanisms in eq. (9) and achieving improved scaling laws through feature learning are particularly interesting accomplishments.

Overall, although this paper only studies a toy learning algorithm, I believe it makes a valuable contribution towards a theory of scaling laws for finite, feature learning neural networks, which is why I recommend acceptance.
summary: This paper studies a minimal model for scaling laws in generalization error under feature learning non-linear neural networks. This model induces improved scaling exponents compared to kernel regression through feature learning (only) for target functions outside of the initial NTK’s RKHS. Consequently, different scaling rules of parameters and training time are asymptotically compute-optimal. These scaling rules are empirically verified in toy settings.
weaknesses: Besides partially critical questions below, I consider the following points to be the main weaknesses of this paper.

- **Artificial learning setup.** The projected gradient descent with frozen first-layer weights is a constructed training procedure. This may make the analysis easier, but it remains unclear how well the results are transferable to more practical optimization procedures like SGD or Adam on **all** weights. Do you think your results can be generalized to realistic training? Do you train all weights in Section 5? This is not explicitly stated.
- **Oversimplifications.** The simplifications that make the paper very easy to read sometimes lead to oversimplifications and imprecision. For example, in l. 191, what is $N,B\gg 1$ formally supposed to mean? Overall, the approach taken in the paper seems very heuristic. For example, starting with finite $M$ and taking the limit $M\to\infty$ omits many mathematical technicalities, why the claimed limiting properties and statements are well-defined and hold.
- A small limitation is that long training times are not well covered (Figure 6b). But this is acceptable as the long training regime is, in general, still very open.

**Typos in lines:** 128, 201, 202, 476.

Meta-review: This paper studies the scaling laws of a solvable two-layer linear model trained by a variant of SGD in two regimes: lazy learning and feature learning (Strictly speaking, three scaling schemes, hard, easy, super easy based on the source condition). Based on the derivation, the paper gives the power-law variation of the loss with training time ( t ), number of parameters ( N ), and batch size ( B ) in these two regimes. The main contribution of this work shows that feature learning can accelerate the convergence of hard tasks. 

The AC also reads this paper. The analysis starts from NTK and then considers the trainable features in Eq. (3) by SGD variants. This scheme indeed simpifies the analysis but is a bit far away from standard neural network training. The key analysis for scaling law used in this paper is based on DMFT, and then the scaling law can be given under souce and capacity conditions. In the AC's view, this paper uses a "tricky" way from the kernel regime to deliver the result for neural networks. But this is still a good paper to the community, so I recommend to accept this paper as a poster. 

Furthermore, the AC suggests the author to carefully polish the paper and re-organize the structure for readability, as well as the discussion about related work and a mathematical definition of feature learning.

Decision: Accept (Spotlight)