How Feature Learning Can Improve Neural Scaling Laws

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?

We are not entirely sure what causes the scaling behavior to change so dramatically on CIFAR-5M. One hypothesis could be that the covariance structure of the task in the limiting kernel eigenspace actually does not follow a perfect power law forever, but has a big gap that cuts off after some large but finite number of dimensions (this could be an artifact of the finite sized generative model that produced this dataset from the original finite CIFAR-10 dataset). The reason we see this in the rich networks is that they optimize over this top set of dimensions more quickly than the lazy model. In principle, we could verify this hypothesis by training a linear/lazy model for much longer on this distribution, but then we could potentially run into overfitting issues related to using a large number of passes over the dataset. This is currently speculation, but we hope to explore what is happening in real and synthetic datasets (like CIFAR-5M) further in future work.

In eq. (6), you state $\lambda_k \sim k^{-\alpha}$ and $\lambda_k^{-\beta} \sim k^{-\alpha\beta}$. Is this a typo?

Yes this is is a typo. We removed the $\lambda_k^{-\beta}$ from this line as it is completely unncessary.

Line 192 , eigenspace

Yes, we will correct this.

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.

We thank the author for appreciating our novel theoretical contributions, despite the limitations of our setting.

Weaknesses

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.

We thank the reviewer for this question and for giving us an opportunity to clarify how our model relates to nonlinear networks. We emphasize that the deep networks we train in section 5 are all trained with every layer updated with SGD with the same learning rate. While our results do not currently try to model training with Adam, we do think capturing training a full nonlinear network with SGD is an advance in realism compared to existing theoretical models.

To see how our model connects to training all layers of a nonlinear network with SGD we first consider the lazy limit of training a deep network $\gamma \to 0$. This is the model of Bordelon et al 2024.

1. In this case, we start with a deep nonlinear network, whose dynamics are governed by its initial (and static) NTK.
2. We start by diagonalizing the infinite width NTK for this architecture which gives a collection of orthogonal eigenfunctions and eigenvalues $\lambda_k$.
3. The finite width kernel eigenfunctions are expanded in the infinite width basis which gives the coefficients $A$ (which will be static over training).
4. The learnable coefficients $w$ are updated with SGD, which will exactly mimic the behavior of the full network will all layers trained in the lazy regime

Now, for the rich regime, the empirical kernel for a deep network evolves over training, which necessitates a dynamical system for $A$. We pick the simplest update rule for $A(t)$ that makes sense with respect to dynamics and model scaling laws (projected-SGD) and recovers the lazy limit as $\gamma \to 0$.

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?

We thank the reviewer for pointing this out. We will try to provide links to additional detailed explanations in the Appendix. We aimed to make the main text of the paper less technical and more accessible while still providing a high level idea of the results.

On the questions of $N,B \gg 1$, what we mean by this is that, while the theory accurately captures the expected loss dynamics, the variability of a randomly drawn system at size $N$ with batchsize $B$ also contains fluctuations around our loss predictions (see the errorbars on Fig 2 c-d or Fig 4). As $N,B$ increase, randomly drawn systems converge to our mean field predictions (dashed black lines). Theoretically, the dynamics of this errorbar in the loss curves could also be computed with additional effort using techniques like those employed here or here which compute fluctuations around a mean field theory.

To address this, we have updated the second footnote on page 4 which now reads as

"There are finite size fluctuations around the mean-field at small $N,B$ which are visible in errorbars in Figure 2 c-d, which could also be extracted from the theory. Alternatively, we can operate in a proportional limit with $N/M, B/M$ approaching constants, which is exact with no finite size fluctuations."

We thank the reviewer for their careful reading and supportive review.

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.

We thank the reviewer for appreciating our attempts to make the paper readable and clarifying the limitations of the current theory so that future research can improve on it.

Weaknesses

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.

While we do focus on the online setting and perform our simulations in this regime, our theory in Appendix D would capture the finite dataset corrections. The $\gamma \to 0$ limit of these were shown to work in Bordelon et al 2024 dynamical scaling law paper (see Figure 2d). While we haven't written a solver for the $\gamma > 0$ data repeat DMFT, we expect the theory to continue to hold under finite data. We try to extract from our theory the data scaling exponent in Appendix E.4. With more time, we could try adding finite data simulations to this paper.

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.

Thank you for this suggestion. While we tried including a bit of detail on what the relevant correlation and response functions are for the reader familiar with DMFT methods, we realize that this may be an unhelpful distraction for readers who are unfamiliar with these techniques. We will therefore remove some unncessary detail from this section to it simpler to parse.

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.

Thanks for this suggestion! We will stress that we are changing coordinates to the eigenbasis of the kernel.

Thank you for the careful reading and comments. We have tried addressing these concerns below.

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).

We thank the reviewer for their careful reading and for appreciating these strengths of our work.

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.

We thank the reviewer for this point. We will revise the draft to provide more intuition for the effect. The main phenomenon is that, for tasks outside of the RKHS of the initial kernel, the empirical kernel will grow as a power law, driving down the loss with an improved rate.

We added some text in the section titled "Accelerating Training in Rich Regime" which states

"This acceleration is caused by the fact that the effective dynamical kernel $K(t)$ diverges as a powerlaw $K(t) \sim t^{1-\chi}$ when $\beta < 1$ (see Appendix \ref{app:bottleneck_scalings}). This is due to the fact that the kernel approximation at finite $\gamma$ is not stable when training on tasks out of the RKHS."

I also feel that the authors do not relate much to the already existing literature 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.

We appreciate the reviewer for pointing us to these relevant works. The feature learning papers that examine lazy/active regimes are related to our $\gamma$ parameter and we will mention this relationship (specifically $\gamma \to 0$ gives a lazy limit like $1/\tau$ in the Kunin et al paper). Specifically, the hidden weights $A$ in our model are updated at a relative rate $\gamma$ compared to the readout weights, which is an unbalanced case considered by Kunin et al. Our dynamics can also be viewed as "mixed" in the sense of Tu et al and their Markovian dynamics closely resemble the dynamics we obtain in the $N \to \infty$ limit (Appendix E.1.2). Lastly, just like the Chizat paper, we also would observe the $1/\text{width}$ convergence rate to the limiting large width dynamics at very early times or in the overparameterized regime when data are repeated. However, our model also captures finite width effects and can exhibit power law model-bottleneck scaling laws at late time which are task-dependent.

Some things that differ between our analysis and some of these prior works are

1. Allows for arbitrary (and as special case, power law) feature structure. We solve for the resulting power law exponents for the loss.
2. Our projected gradient descent can exhibit power law scaling with the width $N$ of the network (not just $1/N$ as is reported in Tu et al), recovering a Chinchilla-like compute optimal scaling law.
3. We also consider the impact of SGD noise and show it only alters the scaling law picture in the super easy task regime, even with feature learning.

We added a paragraph in the Related works section that explicitly references these three works.

"Recent works have examined the dynamics of linear networks, contrasting the dynamics in lazy and feature learning regime, including analysis of infinite width linear networks \cite{chizat2024infinite}, and linear networks with varying and unbalanced initialization and learning rates \cite{kunin2024get, tu2024mixed}. Our model can be interpreted as a two-layer linear network which captures finite width effects (with task-dependent scaling laws) from random initialization. Like these related works, our model also has unbalanced learning rates between hidden and readout weights set by a parameter $\gamma$ that recovers a lazy limit as $\gamma \to 0$."

Strenghts

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.

This paper conducts a series of experiments to verify the theoretical results, making the conclusions of this paper more convincing.

Thank you for appreciating our paper's theoretical and experimental contributions.

Weaknesses

This paper ignores the impact of gamma on the feature learning scaling law.

We actually can characterize the effect of $\gamma$ and will write more about the role of this parameter. While this parameter does not alter the scaling law exponent (and therefore does not alter compute optimal exponents either), it does change the timescale where the loss curve significantly separates from the lazy learning curve. The time when the loss curve transitions from the lazy power law to the new power law is $t\approx \gamma^{- \frac{1}{1- \chi(\beta)}}$ where $\chi = \beta\max\{1,2/(1+\beta)\}$. We note that $\chi < 1$ in the case where feature learning changes the exponent. Therefore, very small $\gamma$ requires a long time to see the transition to the new power law whereas larger $\gamma$ causes an almost immediate transition. At late time, the loss is effectively lower by a multiplicative constant $\gamma^{-\beta}$, ie the loss scales with $t$ and $\gamma$ as $\mathcal L \sim t^{-\chi(\beta)} \gamma^{-\beta}$.

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

Our full theory provides the complete learning curve in terms of $\gamma$ and $t$. We approximate this full expression to get simplied expressions for the exponents (eg in equations 9 and 10). We will comment on the following facts for the hard task case. Approximating our full theory, we can provide a more detailed description of the transition that we will comment on in the paper.

1. For times $t \ll \gamma^{-\frac{1}{1 - \chi(\beta) }}$ the dynamics closely track the lazy learning curve.
2. At timescale $t \approx \gamma^{-\frac{1}{1- \chi(\beta)}}$ the power law transitions from the lazy learning curve to the new power law.
3. For all $t \gg \gamma^{\frac{1}{1-\chi(\beta)}}$, the loss looks like $L \approx \gamma^{-\beta} t^{-\chi(\beta)}$.

Thus we could write something like the following

$

\mathcal L(t) = \begin{cases}
t^{-\beta}  & t < \gamma^{- \frac{1}{\chi - 1}}
\\
t^{-\chi(\beta)} \gamma^{-\beta}  & t > \gamma^{- \frac{1}{\chi - 1}}
\end{cases}

$

We provide this information in a new Appendix section E.2.1 and also reference some of this information in the "Accelerated Training" section.

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

Thanks for this comment. For the final version, we will try to improve the color scheme to make things more differentiated.

From the experimental results in Figures 9 and 12, we can see that the scaling law does not change suddenly at, but changes gradually at a very small value of gamma. Can this paper's analysis cover this aspect?

Yes, we can cover this aspect. Our theory predicts that the loss will noticeably separate from the lazy learning curve at a timescale $t = \gamma^{ - \frac{1}{1 - \chi} }$.

We have tried addressing the concerns and questions raised by this reviewer and hope that in light of these they would consider increasing their score.

Strengths:

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.

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.

Thank you for appreciating these aspects of our work.

Weaknesses

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.

We have tried to improve some of the exposition and motivate the model (see blue text) without relying as much on prior works. We also removed some unnecessary mathematical detail from the main text. We can make additional improvements before the final version of the paper.

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.

We will give a clear operational definition of feature learning: the empirical NTK of the model changes. We state this in the first paragraph of section 2.

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.

This is a fair critique. However, we note that this modeling assumption is basically saying that the empirical NTK can evolve throughout training. The $A(t)$ matrix gives the coefficients for the empirical kernel features at time $t$ in terms of the infinite kernel features at time $t=0$. This matrix exists for deep nonlinear networks, but we do not know apriori how it evolves. This motivated us to consider a simple update rule for $A$.

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.

This is a good question. While our theory does not directly characterize nonlinear networks, we wanted to perform some empirical tests in more realistic settings where we could examine how general our model's predictions were. We suspect that the finite/infinite RKHS norm as a transition point for feature learning since perturbations around the kernel limit are unstable if the target function is outside of the RKHS of the initial kernel (if $\beta<1$). The fact that our new predicted exponent $2\beta/(1+\beta)$ still seems to work on nonlinear tasks is fortunate, but not guaranteed by our theory and is currently conjectural.

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

Thank you for pointing this out. We will try stressing the main points about feature learning which are

1. Feature learning may only change scaling laws for tasks which are sufficiently hard for the initial kernel.
2. In our model with natural power law covariate structures with the target function outside the RKHS, feature learning improves the scaling exponents for optimization but not with respect to model size.
3. The scaling law for SGD in this model with respect to training time improves from $t^{-\beta}$ to $t^{-2\beta/(1+\beta)}$ while the parameter scaling laws do not change.

We added the following sentence below our listed contributions

"Overall our results suggest that feature learning may improve scaling law exponents by changing the optimization trajectory for tasks that are hard for the initial kernel."

Thank you for providing additional clarifications and for the efforts to expand the discussion in Appendix H. I appreciate the engagement and the inclusion of definition B and task alignment, which indeed offer a meaningful direction to explore the concept of feature learning further.

The two definitions we provided previously are indeed compatible with each other..."

I did not mean to suggest that the definitions were contradictory. However, I maintain that they are distinct in focus and scope. The definition of feature learning as "the empirical NTK of the model changes" seems overly broad and potentially misleading. While the second definition, which incorporates the lazy regime and optimization perspectives, offers a more refined approach, it still falls short in directly addressing the quality or utility of the features learned. In particular, observing changes in the loss function alone does not provide sufficient evidence to conclude that meaningful or task-relevant features are being learned.

That said, the added Appendix H is a promising and commendable addition. Specifically, definition B and the concept of task alignment resonate well with the broader understanding of feature learning. I would be particularly interested in seeing how alignment evolves across tasks of varying hardness, especially for tasks with β<1. Unfortunately, the main plots in the paper predominantly focus on the loss function, which does not provide much insight into the dynamics or utility of feature learning.

I recognize the contribution of the paper in demonstrating the interplay between task hardness and the lazy or rich regime scaling law, and Table 1 serves as an excellent summary in this regard. However, the framing around feature learning remains problematic. The results and analyses in the paper could be entirely framed around the lazy and rich regimes without invoking the concept of feature learning.

In light of this, I would strongly recommend either (1) revising the framing of the paper to remove claims about feature learning, or (2) building upon definition B to empirically study alignment changes across different tasks. Without either of these adjustments, the connection between the presented results and feature learning remains tenuous and detracts from the otherwise valuable contributions of the paper.

We do not claim that feature learning always improves scaling laws...

This is an important point, but I think what this paper shows is that being in the rich regime does not always help scaling laws, and indeed the lazy regime might be just as good. This is probably related to the feature learning and alignment argument, but with the current results, I am not sure what we can say about feature learning in the sense of definition B in Appendix H.

Final word: I understand that it is too late to make significant changes, but it would be great to see the alignment evolution explored further using definition B. That said, I strongly suggest adding definition A earlier in the paper for clarity. Based on definition A, feature learning is simply defined as being in the rich regime and allowing the weights to change. While I do not believe this definition is reasonable, I understand that it has unfortunately been repeatedly used in prior literature. In light of this, I have decided to increase my score to 6.