Grokking in Linear Estimators -- A Solvable Model that Groks without Understanding

Thank you for your thoughtful review and constructing comments and questions. We are glad that you acknowledge the main goal of our work, which is to call for a more serious discussion on what grokking is, and properly define it by multiple metrics that must converge to a clear conclusion.

Before addressing the weaknesses raised in the review, we must first provide a small correction to how our results were perceived. Most of our analysis was done in the regime $d_{in}/N_{tr}<1$, meaning that there are more training samples than input dimensions. In this case, there are no zero eigenvalues, and the difference between training and generalization loss comes from the mismatch between the covariance matrices for the nonzero eigenvalues.

Next, we address the weaknesses in the order in which they were raised.

1. Is this grokking?

This is indeed an important point that also came up in the reviews of other referees. We agree, and it is in fact one of the main conclusions of our work. Namely: we study a model which does not exhibit any feature learning, no special weight structure or non-monotonicity in its generalization loss evolution, and can still be considered to "Grok" if one uses accuracy as the only metric (and this specific case was actually descried as grokking in the Omnigrok paper [Liu 2023]). Truly, if this nearly trivial model Groks, then any model can Grok.

We believe that this aspect of delayed generalization should be separated from representation learning grokking which is present in other systems, such as modular arithmetic. However, we believe that some aspects of the proposed mechanism -- namely that the accuracy is simply a step-like-function of the loss that depends on the data covariance -- might be quite general. We provide some preliminary evidence in the new appendix G.

We agree with the suggestion of tracking both generalization loss and accuracy to define grokking, but even in this scenario there may be other factors, such as optimizer dependence that can induce delayed generalization on the loss, without any deep change in learning behavior.

We completely rewrote tyhe discussion to better emphasize this point

2. Similarity to previous works

While we recognize some similarities to previous work, we believe our study makes several novel contributions:

We provide a closed form solution to the dynamics of training for linear models, rephrasing previous asymptotic analyses in simple Random Matrix Thoery language. This allows us to characterize training dynamics over the full training process and isolate the contribution of different aspects.

By relating training loss directly to the spectrum of the data covariance matrix, we suggest that perhaps some of our results can be extended beyond the linear regime, as shown by the 2-layer $\tanh$ example, as well as a new appendix (Appendix G.) in which we examine the dependence of the grokking time on $\lambda$ for a modular arithmetic task.

3. Finite time analysis vs previous works: We agree that previous works have examined the exponential decay of the discrepancy in linear models, and therefore the novelty in our work is relating that same discrepancy with the phenomena of grokking. We added some of your suggestions to the bibliography.

We hope with these changes the main message of the paper is clearer, while the broader context is more readily accessible, and that upon review of the revised version you find our work suitable to acceptance.

Thank you for your thoughtful review and positive appraisal of our work. We are glad that you found our analysis of grokking through a theoretical lens in linear models and neural networks to be a novel contribution that substantially advances understanding. You also highlighted how our comprehensive theoretical predictions and extensive empirical validations effectively supported our claims. These are certainly important strengths of the work.

At the same time, you raise valid points about areas we can improve.

1. Notations: We agree. We edited the manuscript to make the notations clearer and easier to follow. Thank you for your comments.

2. Theorems: Following your comments, we have added a short introduction to the model at hand, clarifying the notations we used. We have also added references to certain theorems which substantiate our asymptotic results in more limiting cases.

3. NTK limit: Thank you for the question. We believe that the answer is that in the lazy regime, i.e. when the kernel remains fixed at initialization $h_t\simeq h_0$ our analysis will hold for NTK, however any deviation from the strict NTK limit may generate dynamics that differ for our results, in any truly nonlinear settings. This implies that some of our predictions regarding the scaling of grokking time with certain parameters, for instance, $d_{in}$ and $N_{training}$ should hold, at least to some extent, until breaking away fully from linear dynamics.

Thank you for your careful reading and thoughtful comments. We are glad that you believe we presented a sound and comprehensive analysis of our simple model which captures some aspects of grokking.

In light of the weaknesses and questions brought up, we have made a number of modifications to the text, as detailed below.

• Comparison to gradient flow in linear models literature and related works: We apologize for the errors found in the bibliography and the related work section. Particularly, the work of Gromov was one of the motivations for our work and was somehow unintentionally missed. We thank the reviewer for bringing these highly relevant works on linear models (for instance Saxe et al.) and on gradient flow to our attention. they are now incorporated into the main text.

• Extension of the results to more realistic setups: We fully agree that our model does not capture all aspects of grokking. As we wrote in our response to other referees, this is actually a main point of our paper: in some cases grokking is just an artifact of the accuracy measure, and does not result from an interesting thing that goes on in the underlying dynamics. This was stated explicitly in the abstract and the text. We rewrote the discussion to stress this point.

We agree with your ultimate conclusion: "...the primary point of the paper is motivating the need for more careful research on grokking in the future. For example, what artifacts due to choice metric should experimenters be wary of? How should we distinguish “improvement in understanding” vs mere statistical noise?..." . This was a main goal of this work. Since grokking is ill-defined, even a linear model which does not "understand" anything deep, can still "grok" because of an artifact.

That said, it is still possible that at least some aspects of our analysis would extend to realistic setups, for instance the grokking time dependence on $d/N$, at least for some models. That is, at least some of what is currently considered "grokking" can be explained using our formalism. We begin to address this point in the new appendix G, but this is clearly an interesting point for future research.

We stress in the discussion that our work is not intended to be a "theory of grokking" in general.

Regarding the specific aspects of grokking that are missing in our model:

1. Specialized weight structure: Since the linear model does not exhibit any type of feature learning or gains any deep insight during training, there is no specialized structure for the weights. Indeed, this is a deviation from grokking in modular arithmetic.

2. The effect of weight decay: In the linear model, the effect of WD can be analytically understood, and in fact the reason for the difference is that they explore WD values in our "gray" regime, as our model cannot generalize in that parameter space while their setup can. We have revised our manuscript to explain this point.

3. Loss increase prior to grokking: Indeed, this is also a feature that does not exist in the linear model, as all the quantities in the model change monotonically with one another, and so unless one appeals to a catapult scale learning rate in the two-layer example, this cannot happen. However, our preliminary results (in the new appendix G) hint that this non-monotonicity might not be crucial for grokking (see figure 8, right panel, where the non monotonicity is present).

Questions/Comments:

1. $\tanh$ gain: Yes, we believe that the near linear behavior of $\tanh$ is the reason for the matching results, since the analysis would not hold for a general activation function. For instance, we find empirically, that using ReLU activations completely deviates from the simple prediction of $h_t\simeq h_0$. Extending $g>1$ does break the predictions as well, above a certain threshold. Due to time constraints we were not able to incorporate a discussion in our work.

2. Averaging over teachers: It is probably true that the average over teachers would turn approximations to equalities. We show a weaker analytical claim, but a stronger empirical verification: our model give effectively perfect predictions for the dynamics of individual trajectories, and not average predicitons on ensembles.

3. Teacher decomposition on the kernel This question requires more substantial analysis - while it is true that the decomposing the teacher in the eigenbasis of the kernel would imply dependence on the eigenfunctions themselves, the question of how these functions evolve during training would likely determine their importance. In the infinite width limit, since the kernel is fixed this effect should not exist, but we agree that there are some non-trivial calculations to be made.

Thank you for your positive feedback on our paper. We appreciate you highlighting its clarity and its thorough explanations of the theoretical analysis and empirical results through a tractable model. Below we address the weaknesses and questions raised, in hopes of strengthening the message of our work.

Weaknesses:

We agree with your concerns that our conclusions might not extend beyond the linear model, and that our model does not capture some of the phenomena associated with grokking, e.g. a non monotonic loss. A main goal of this work was to highlight that the grokking phenomenon is not sufficiently well defined, and may contain multiple aspects, some of which can be captured (or mimcked?) by a nearly trivial model which shows no feature learning at all. This is indeed the main point we are trying to make, as we hint in the title of the manuscript.

It should also be noted the model was not introduced in this work - we are analyzing a grokking scenario reported in the Omnigrok paper, and show that grokking in this nonlinear case can be understood in terms of a linear model (as we show in the two layer $\tanh$ example). That is, at least some of our analytical predictions might extend to nonlinear models. For instance, the change in grokking time as a function of $d_{in}/N_{training}$ is expected to hold in some form for more realistic setups as well. We have added experimental evidence in the revised manuscript that indicate that the grokking time scaling behavior extends for deeper nonlinear networks, for the task of learning modular arithmetic. We hope this change will provide a sufficiently broaden the scope of our work.

Questions:

1. Delayed generalization vs. grokking: This is an excellent question, which is related to the previous point. Indeed, we agree that if grokking is defined by representation learning - hence requiring a dramatic change both in loss and in accuracy, this model does not grok but just pretends to. For instance, grokking in modular arithmetic clearly displays an ascent of the test loss while training accuracy saturates, and a drop of the test loss when the test accuracy increases from 0 to $100\%$. One of our goals in this work is to precisely separate the effects which depend only on the accuracy measure, from true representation learning. We also have preliminary results which may indicate that even the loss behavior is not necessarily indicative of "representation grokking", which we plan to present in future work.

2. Self correction and label noise: It is true that label noise was only commented upon in a paragraph in the original manuscript, but we can easily explain its effect here. Assume training in the most basic setup, in which $d_{out}=1$ so $S,T$ are vectors and there is a scalar label noise variable $n\sim\mathcal{N}(0,\sigma_n^2)$. Training with label noise in our model can be shown to give the average training loss $\mathcal{L}_{tr} =Tr(D_0 D_0^T e^{-4\eta_0 \Sigma_t } \Sigma_t )+\sigma^2_n$, while the generalization loss (w/o noise) is unchanged, thus adding noise to test labels will add a similar additive term. Therefore, the training dynamics do not change except for an additive constant. It is then possible to tune $n$ and the threshold parameter $\epsilon$, such that the noise can cause induce grokking precisely when test accuracy saturates. There is no sense of self correction here, just a tuning between noise and the accuracy metric, so it is unclear if this is the same mechanism at work in modular addition. We thank you for highlighting this point and we included an appendix explaining this.

Figures: We thank you for this comment, and we have reworked our figures to be more clearer and more legible.