Paper: Grokking as the transition from lazy to rich training dynamics (ICLR-2024)

Reviewer comment (Theory):
One point I think worth being clarified is the role of weight decay. I think we have a pretty good understanding of how weight decay changes the kernel. Essentially, weight decay almost always changes the NTK at the speed of $\kappa$, which is the strength of weight decay. The time scale for the effect of weight decay to take place is nothing but $1/\kappa$. Therefore, if $1/\kappa$ is far larger than the time scale of lazy training, one has a separation of time scale, where lazy training happens first, and feature learning induced by weight decay happens later. I think the authors should make more discussion on this point as it constructively adds to the paper.

Author response:
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.

Reviewer comment (Theory):
The theoretical explanation presented in this paper only works for the new definition of grokking, but not the original definition of grokking in Power et al., 21. For most experiments in this paper where the authors claim grokking happens, two criteria of the original definition of grokking are not met: 1. when test loss starts to decrease before training loss gets small; 2. the eventual test loss is still high. Figure 5(b) is an example.

Author response:
**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.

Reviewer comment (Theory):
There is no end-to-end theorem about grokking. Given the modified definition of grokking, explanation using transition from kernel regime to feature learning regime seems quite straightforward. It should be very standard to show that there exists a period of time where training loss can decrease but test loss has a lower bound due to living in the kernel regime. Still, this paper lacks such precise and rigorous statement.

Author response:
We thank the reviewer for the question. 

**Revised and new sections now in blue**

We have changed the color of the modified sections (Section 5.3, Appendix Section 13, Appendix Sections 14-18) to blue to make it clearer what the new additions are. 

**What is DMFT?**

Dynamical mean field theory is a method for obtaining the train and test loss dynamics for high dimensional data and large datasets. The theory operates in the limit of $P,D \to \infty$ with $P = \gamma D$. The summary statistics in equation 4 concentrate in this limit and characterize the train and test loss dynamics, from which we can predict grokking behavior.

**The need for dynamics: motivating DMFT**

Thank you for your theoretical question about timesteps, it cuts to the heart of the issue. The two conditions you ask from a theory of grokking can be restated as saying a theory of grokking must demonstrate there exist $t_1<t_2$ so that: 
1. $L_{test}(t_1)-L_{train}(t_1) \gg 0$ is large with $L_{\text{train}}(t_1) \approx 0$ small (a time at which we memorize but do not generalize). 
2. $L_{\text{test}}(t_2) \approx L_{\text{train}}(t_2) \approx 0$ is small (end-time generalization). Notice to show this it suffices to show $L_{\text{test}}(t_2)-L_{\text{train}}(t_2) \approx 0$ and $L_{\text{test}}(t_2) \approx 0$. 

The fact that we want to reason about this test-train gap *at different points in time* means that it does not suffice to prove facts about *static* quantities like lower bounds on end-time test loss, etc, and that we must reason about various quantities (weights, activations, this test-train gap) *dynamically*. DMFT is just a common technique [(Mignacco et al., 2020; Celentano et al., 2021)] to do exactly this. Note we can only use this mathematically involved technique to reason this way because our setting is as simple as it is: doing theory on dynamics for more complicated models like one-layer Transformers (Nanda et al, 2023) or MLPs with AdamW (Liu et al, 2022b) is untenable, for instance. 

**Figure 4 exactly demonstrates the existence of $t_1, t_2$ is predicted by our theory**

Consider Equation (7): 

$$L_{\text {test }}(t)-L_{\text {train }}(t) \sim \frac{2 \alpha}{\gamma N} \sum_{i j=1}^N \int_0^t d s R_{i j}^w(t, s) C_{i j}^g(t, s)+\mathcal{O}\left(\gamma^{-2}\right)$$

From this, we can make observations about $t_1, t_2$ at two levels of abstraction

- See how the test-train gap dynamics scale qualitatively with $\gamma, \alpha, \epsilon$, to predict what tuning of these parameters gives grokking. On page 8, the "Overfitting effects at finite $\gamma$" and "Effect of $\epsilon$ and $\alpha$" do exactly this, and experiments (eg. Figure 3) to verify these insights. 
- **DMFT outputs implicit integral equations like (7) that can be solved numerically. We want to emphasize the DMFT plots in Figure 4 are the predictions of our theory, so the existence of a test-train gap in those curves obeying conditions (1, 2) from earlier means the required $t_1, t_2$ exist and can be read off the curves. For instance, Figure 4(a)(right) shows how $\gamma = 15$ has $t_1 = 20, t_2 = 50$ where initially train loss halves before test loss starts moving at all. This would become even more emphatic for smaller $\gamma$ (ie. train loss would go near zero before test loss starts moving -- exactly what is demanded in conditions (1,2)), but solving the DMFT integral equations is numerically challenging there. Similarly, in Figure 4(b)(right) for $\alpha = 10^{-1}$, we see $t_1 = 18, t_2 = 60$ satisfies the two conditions also.**

The reviewer might be wondering about simple, closed-forms for $t_1, t_2$. We address this next. 

**Cannot get $t_1, t_2$ in simple closed form because train and test loss dynamics are computationally hard**

We cannot get $t_1, t_2$ in simple closed-form, *because knowing test and train loss at every $t$ is, in some sense, at least as computationally hard as running gradient descent.* The key is that DMFT outputs *implicit* equations characterizing loss over time, hence the need to solve them numerically. More specifically, it is hard to calculate simple closed-forms for the test-train gap because because the final answer depends on solutions to self-consistent equations for the summary statistics defined in (4), which are computed as averages over non-Gaussian stochastic processes (5-6). Thus it is challenging to "read-off" timescale gaps simply by inspecting the integral equations alone (without solving them numerically). 

We thank the reviewer for their question and hope this illustrates how our theory of grokking does indeed predict such a $t_1 < t_2$ when its final equations are (numerically) solved out, as desired. The key Figures illustrating this are Figure (4) and Figure (13).

Reviewer comment (Novelty):
As mentioned in the strengths, the explanation of grokking using transition from kernel to feature learning regime (Li et al., 2021) is known even before the original grokking paper (Power et al., 2022). Though this paper is novel in the sense that it focuses on vanilla GD, which is more natural than Li et al., 2021, the theoretical result in this paper is also not as clean as that in Li et al., 2021.

Author response:
**Comparison to Li et al**

While we appreciate the contributions of Li et al on the implicit bias of SGD dynamics with label noise, and on lower bounds for kernel generalization error, we point out the following novel contributions in our work. First, our results and experiments do not rely on either SGD or on label noise, so explanations which rely on these phenomena are insufficient to explain our setting.  We also believe our conceptual contributions go beyond that analysis in this new setting studying grokking
1. Studying how rapid feature learning can reduce the gap in time between train loss and test loss decreases
2. Showing how harder tasks for the initial kernel lead to more extreme grokking. 
3. Providing an asymptotic description of networks (their weights, activations, train and test loss) trained on random data *far from the kernel regime* using DMFT techniques. We have provide an approximate equation for the test and train loss gap valid for large but finite $\gamma = P/D$ (see Equation 7).

We hope in light of these improvements and contributions, the reviewer will consider increasing their score. We thank you for your detailed comments.

Reviewer comment (Experiments):
(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 spans more than one direction) play a role in the presence of grokking?   
(Nichani et al. 2022) *Identifying good directions to escape the NTK regime and efficiently learn low-degree plus sparse polynomials.*

Author response:
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.

Reviewer comment (Theory):
The target function is a single-index model with information exponent 2 (Ben Arous et al. 2021). It is known that for such a 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?   
(Ben Arous et al. 2021) *Online stochastic gradient descent on non-convex losses from high-dimensional inference.*

Author response:
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.

Reviewer comment (Novelty):
The quantity describing the difficulty of target function \(y^\top K^{-1} y\) has been extensively studied in the kernel literature. See (Arora et al. 2019) for the analysis of NTK.   
(Arora et al. 2019) *Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.*

Author response:
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.

Reviewer comment (Theory):
(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.

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