Grokking as the transition from lazy to rich training dynamics

We thank their reviewer for their comments and suggestions to improve the theoretical components of the paper. Below we discuss their questions and weaknesses and describe what we did to improve it.

Limited Theoretical Component to the Work

Asking for more theory was common amongst the reviewers, so, as outlined in the global response, we have included theory we had not included in the initial submission. Specifically, we derive limiting dynamics for quantities that determine the train and test losses when two layer networks are trained on high dimensional data. We show that the derived theory is accurate in Figures 4 and 13 and use it to derive approximate test loss and train loss gaps as functions of training time. Because these results are now so much stronger than the loss decomposition, we moved that to the Appendix.

Can Grokking Happen if Rich Regime Can't Generalize?

You are right that since grokking is roughly "delayed generalization," the network must be able to generalize in the first place from this sized training set. You make a good point here that some tasks may not even be soluble by a given network architecture in the rich regime no matter how much data you give it, because, for instance, the network architecture may not be expressive enough or have a high enough model capacity. To be clear, you are right that we assume the network is capable of generalizing on the task to begin with to be able to see grokking.

NTK Alignment and Data Quantity

This is an excellent question; we have added a new Section 17 to the Appendix to test this explicitly. We show that indeed even for task labels set to small eigenvectors, large amounts of data suffices to learn them and generalize. However, if you let the task become a function of the dataset size by defining the labels to be the smallest eigenvector of the empirical kernel matrix (in eigenvalue), then the task gets commensurately harder as you increase dataset size, and so even feature learning cannot allow this architecture to learn the task. This corresponds to going down the columns of our new Figure 21 and seeing that if we set $j = P$ then no amount of data will allow this architecture to generalize perfectly, giving no grokking.

We thank the reviewer for their supportive comments and careful reading. Below we attempt to address their question about weight decay.

Role of Weight Decay

You are absolutely right. As you point out, papers like (Lewkowycz et al, 2020) show that the NTK evolves linearly as a function of weight decay. In response to your review, we have reworked Section 14 of the Appendix to explicitly discuss the role of weight decay, finding that past observations (such as those by Nanda et al, 2023, and Liu et al, 2022b) that empirically find weight decay has an approximately linear relationship with time to generalize (grok), our theory of grokking as a transition from lazy to rich dynamics predicts this exactly because, as you point out, the NTK evolves at $O(\kappa)$, so that the network is forced to leave the lazy (static NTK) regime at $t \sim O(\kappa)$ as well, which is what has been observed in past papers like Liu et al 2022b, but without explanation. Section 14 also illustrates how the laziness parameter $\alpha$ can supplant and replace weight decay in this way: making a network richer by decreasing $\alpha$ can serve the same effect as forcing NTK evolution (feature learning) by increasing weight decay, $\kappa$, and we showcase an example of that trade-off in Figure 14.

Difference Between Accuracy and Loss Curves

The difference between accuracy and loss curves is important to clarify, so thank you for asking this.

First, we added a two-page Section 17 to the Appendix explicitly addressing the questions here, including four new experiments to support our claim that loss is the right primitive to study. But we begin by first addressing your main concern that accuracy based curves in (Power et al, 2022) are more surprising than the loss curves we study in this paper because of the two criteria you give: (Power et al, 2022) give the loss curves corresponding to their striking accuracy curves (Fig 1) in the appendix of their paper, as Figure 4. They look like the loss curves we study in our paper here! The loss in our main toy model of study (Fig 1, Fig 3, Fig 11, Fig 12) looks exactly the same, and even in other experiments we do with more complicated data (Fig 6b, for instance), our curves also look like that. As far as your citation of Fig (5b) is concerned, we created a new Appendix Section 17 to show how curves that reach nonzero (but low) end-time loss give rise to accuracy curves that exhibit striking grokking. Note also that (Gromov, 2023) finds grokking in a modular arithmetic task with curves that look very much like (5b) as well.

In the new Appendix Section 17, we reproduce grokking on a modular arithmetic task (the task in Power et al, 2022) and show the striking accuracy curves we associate with grokking in Figure 20 and note the following facts

• The test loss is decreasing before train loss is zero
• The test loss saturates to a nonzero value at late time while the test accuracy reaches $100\%$.

This shows that the types of loss curves we document in our experiments are actually the same kind of loss curves as those in original grokking experiments.

We thus disagree with the claim that the original accuracy-based result from Power et al. is stronger (more surprising) than the corresponding loss curves that we treat here, and in fact they are equivalent. As we expect, our conclusions about grokking in loss naturally carry over to grokking in accuracy (but not vice versa). This is why when we intervene on classification-first tasks like modular arithmetic (Fig 2, Fig 6b) and MNIST image classification (Fig 6a), we see our interventions control grokking exactly as we predicted. If our definition (reasoning about loss) was a weaker one, we would not see this in classification tasks, where accuracy is plotted.

We end by noting that we are not the first to make this observation. Indeed, the fact that loss, not accuracy, is the interesting thing to study was quickly realized after the initial paper discovering grokking was published! It is noted when (Davies et al, 2023) study loss in their Figure 5b, when (Liu et al, 2022b) characterize the loss landscapes of various tasks that exhibit grokking, when (Nanda et al, 2023) plots excluded and restricted loss as progress measures in and remarks that accuracy can be gamed, and strikingly in the recent (Schaeffer et al, 2023) that shows how accuracy and other metrics that are not optimized by the network can be very misleading to plot and easily gamed. Thus our choice to do so in this paper is consistent with past literature on grokking. We refer you to Appendix Section 17 created specifically to clarify this perceptive question you raise.

End-to-End Theory of Grokking

While the original paper did not attempt to provide an end-to-end theoretical description of the grokking dynamics, we now include a dynamical mean field theory (see the new Section 5.3 and Appendix 13) for networks trained on high dimensional data. In particular, we can write down simple expressions for the gap between test and train losses as functions of time (see equation 10) which states the test train loss gap at large $\gamma$

$\mathcal L_{test} - \mathcal L_{train} \sim \frac{2\alpha}{N \gamma} \sum_{ij=1}^N \int_0^t ds R^w_{ij}(t,s) C_{ij}^g(t,s) + \mathcal{O}(\gamma^{-2})$ where $R^w$ and $C^g$ are certain functions which can be computed from our theory. The test and train loss dynamics at large width in the lazy regime should emerge as a special case of this more general theory.

First Paragraph of Introduction Sounds Slightly Wrong

You make good points. We have clarified what we meant and rephrased the introductory paragraph. We now elaborate on the points we makes to set up the background motivation for the paper, notably with gentler language than before.

We thank the reviewer very much for their detailed review. Before we address the questions posed in the weaknesses section, we wanted to address some comments in the "Strengths" to make sure we are on the same page about the grokking literature and how this paper situates itself therein.

Strengths: ...This paper proposes that grokking can be triggered by the transition from kernel regime to feature learning regime. Though this is already known even before the grokking paper by Power et al, 2022, e.g., the example of quadratically overparametrized linear model in Section 6 of Li et al., 2021, the novelty here is that this paper focuses on vanilla GD. In contrast, the transition between kernel and feature learning regime, Li et al., 2021, is triggered by sharpness minimization implicit bias of label noise SGD...

First, you are correct that the main proposition here is that grokking can be seen as the transition from a kernel to feature learning regime. And to be clear, we are not claiming to be the first to show a network can transition from a lazy to rich training regime. Our novel contribution is to use this transition as a possible explanation of the grokking phenomenon, which to our knowledge has not been done before.

1. First, we thank you for bringing the paper by Li et al (2021) to our attention; we have cited it as it is relevant. However, it is a paper about the implicit biases of SGD around local minima, and how SGD can trigger a transition from lazy to rich training dynamics. However, SGD is neither necessary nor sufficient to see grokking, so no complete theory of grokking can be one entirely about SGD dynamics.
2. None of the existing papers on grokking mention the notion of NTK or lazy and rich training paradigms (including [(Liu et al, 2022a), (Nanda et al, 2023), (Liu et al, 2022b), (Davies et al, 2023), (Thilak et al, 2022)]). It is true that many papers remark something along the lines of "feature learning is somehow important" and "there is something resembling a phase transition happening" but the key contribution of this paper is precisely in formalizing that claim in a way that allows us to demonstrably control and manipulate grokking in a way nobody has been able to do before.
3. Many of the current popular theories of grokking claim that it is fundamentally due to weight decay [(Nanda et al, 2023), (Varma et al, 2023)] or weight norm [(Liu et al, 2022b)] or adaptive optimizers [(Thilak et al, 2022)], or any number of other things that are directly inconsistent with the fact that grokking arises when a network moves out of the NTK regime (as we discuss more in Section 2 and Section 10 of our paper). We think our paper can help illustrate how mechanically lazy/rich training regimes can alter the grokking curve.

3. The target function is a single-index model with information exponent 2 (Ben Arous et al. 2021). It is known that for such target function, gradient descent on the population loss will start from an approximate saddle due to the random initialization. Is this observation related to the test loss plateau in grokking? And do we expect similar findings when the information exponent is larger than 2?

Indeed, we believe these phenomena are related. In Appendix Figure 10, we exactly tested the effect of "warm start" initial conditions where the weights are partially aligned to the target direction. All things equal, the models with poor alignment (close to saddle point of dynamics) exhibit more dramatic grokking behavior. Conversely, weight initializations with some component of the solution weights to begin with (far from saddle point) have train and test curves that track each other more strongly over the course of training, not exhibiting grokking to the same degree.

You also make a good point about the fact that the setting with information exponent $k=2$ is perhaps an easier task (in the sense of escaping the saddle point at initialization) than the case where $k \ge 3$, as (Ben Arous et al, 2021) outline in Theorem 1.3 and 1.4. Thus, it is natural to ask whether the grokking that arises does so because of any properties related to the quadratic target. We can test this by seeing if grokking arises for $k \ge 3$ higher degree Hermite polynomials. We find that indeed it does, out-of-the-box, so we can be confident that the findings do not rely on special properties of the quadratic target (information exponent $k=2$). We have new experiments illustrating this in Section 15.2 of the Appendix.

4. The quantity describing the difficulty of target function $y^T K^{-1} y$ has been extensively studied in the kernel literature. See (Arora et al. 2019) for the analysis of NTK. (Arora et al)

We have updated the paper to emphasize that this is not a novel metric we are proposing, but rather is a natural, well-studied, object which reflects the RKHS norm of the targets on the train distribution, and turns out to be useful in the examination of grokking. We now cite Arora et al and other appropriate papers when introducing this difficulty metric.

5. (minor) The cited references on the quadratic sample complexity lower bound for kernel estimators such as (Ghorbani et al. 2019) do not rigorously cover the case of Gaussian data.

This is a good point, we have updated the paper to qualify these references when we point to them.

Thank you for the detailed comments! We believe the additional theory (DMFT for dynamics) and experiments (new Appendix Sections 15-18) substantially improved the paper as now the emergence of a test-train gap is theoretically justified in addition to empirically controlled by the parameters we introduce. We hope that the reviewer will consider raising their score based on the new theoretical contributions, and additional Appendix sections.

We thank the reviewer for their indepth reading and their useful suggestions. We took their suggestion to provide a more predictive theoretical description seriously and added a whole section dedicated to this. We also investigated multi-index and higher degree hermite polynomial target funtions.

Below we respond to each of the weaknesses.

Weaknesses:

1. Given the idealized setting (Gaussian data with identity covariance, single-index target), it is rather underwhelming that the authors did not provide any quantitative characterization of the training dynamics to prove the existence of grokking. Instead, the proposed explanation is only verified empirically, which limits the contribution. Can the authors comment on the technical challenges in analyzing the gradient flow trajectory for this quadratic model?

We had begun deriving the dynamics of the test risk such a model, but did not include it in the initial submission because solving the resulting equations numerically was challenging. However, the reviewer's questions about theory have prompted us to complete and include this theoretical work that can potentially shed light on where grokking comes from in terms of data and richness. We believe this substantially strengthens the submission.

In particular, we now derive a dynamical mean field theory (DMFT) description of the dynamics in the proportional regime where $P = \gamma D$ with $P ,D \to \infty$ (with some subtleties about initial conditions for the alignment of weights to target direction). See Section 5.3, Appendix 13 and Figure 4.

This results in a version of the classic Saad-Solla [1] equations for two-layer network dynamics, but with data repetition which induces correlations in weight updates and preactivations across time (allowing for the possibility of overfitting/grokking).

The key mathematical challenges that arise from having a separate train/test distribution when fitting a nonlinear function in this high dimensional limit are

1. Non-markovian dynamics for the weights and preactivations (see equations 7-8), which give compouned errors between the online ($\gamma = \infty$ dynamics) and offline ($\gamma$ finite) solutions. (For example, see the matrix $R^h$ in Figure 13d which is no longer diagonal!)
2. Non-Gaussian preactivation statistics (see equation 6).

However, we can still recover interesting dependences on $\alpha$, $\gamma$, and $\epsilon$ (see Figure 4). For example, reducing $\alpha$ and increasing $\gamma$ reduces the gap between train and test loss dynamics.

$

\mathcal L_{test}(t) - \mathcal L_{train}(t) \sim \frac{2 \alpha}{\gamma N} \sum_{ij=1}^N \int_0^t ds R^w_{ij}(t,s) C^g_{ij}(t,s) + \mathcal{O}(\gamma^{-2})

$

where $C^g$ is the gradient correlation function and $R^w$ is the response function for the weight dynamics. All of these functions implicitly depend on $\gamma, \epsilon,\alpha$. The DMFT allows us to recover from a theoretical framework, the phenomena we observed empirically in our polynomial experiments.

[1] https://journals.aps.org/pre/abstract/10.1103/PhysRevE.52.4225

2. (Nichani et al. 2022) also studied the transition from the NTK to the feature learning regime, and established improved sample complexity for learning low-degree + sparse high-degree polynomials. For the current model, does the sparsity of the higher-order term (that is, whether the quadratic component of the target function span more than one direction) play a role in the presence of grokking?

There does not seem to be a major change in the grokking dynamics when we fit multi-index functions compared to single-index functions. We've added a new set of experiments in Section 15.1 of the Appendix verifing that the grokking we present holds out-of-the-box on multi-index models where the quadratic component spans multiple directions, illustrating that sparsity is not necessary for our results to hold.

We thank you for bringing (Nichani et al, 2022) to our attention. It is a very pertinent example of a transition from lazy to rich dynamics and what kinds of feature learning corrections improve generalization when transitioning away from the NTK. We added a citation to this work.