Paper: Grokking in Linear Estimators -- A Solvable Model that Groks without Understanding (ICLR-2024)

Reviewer comment (Theory):
Grokking is a somewhat ill-defined phenomenon. However, like many grey area situations, the fact that it is difficult to draw one sharp reason doesn't mean there cannot be a sharp feeling that something does fit. The reviewer provides several reasons for why they don't think this is grokking: (A) It is too simple — if this model Groks, everything Groks. (B) There aren't any feature learning effects (see for instance Gromov (2023) or Misra (2022)) or beyond NTK effects. (C) There is only one time scale here, which is (1-\sqrt{\Lambda})^{-2}, governing both test and train loss. Hence, looking at the train and validation loss, a practitioner would not be encouraged to do any early stopping here. In contrast, in various Grokking models, there is an over-fitting regime prior to Grokking which can mislead practitioners.

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

Reviewer comment (Novelty):
Technical aspects seem quite close to previous results in the literature. The training dynamic the authors find is closely related to the e^{-\Theta t} \Delta obtained in the NTK work. In fact, up to a standard shift of the inner and outer indices, their data covariance matrix is Theta here, and in particular both have the same non-zero spectrum. The reviewer is certain there are much earlier works which solve this for a linear network.

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

Reviewer comment (Novelty):
Similar factors relating train and test performances, however at infinite time, were also found in Adlahm and Pennington 2020 and https://arxiv.org/abs/2006.09796. Given the expected exponential decay of the discrepancy in linear models, the fact that this ratio holds also in the asymptotic dynamics seems expected. The reviewer suggests it would be good to compare and disentangle the factors the authors find from those.

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

Reviewer comment (Writing):
The notations are quite dense even for the simplest model at first glance. The paper would benefit from having some explicit definition or intro to the model and all the notations.

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

Reviewer comment (Theory):
The derivation of this paper is concrete. However, all the calculations are in the asymptotic manner. It would be better to see the results be displayed as some rigorous theorem in some limiting case.

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

Reviewer comment (Theory):
In this linear model, the grokking phenomenon is due to the gap of convergence speed between training and generalization loss. Can it also imply the difference in two-layer networks settings (or in NTK regime), or they have a different mechanism?

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

Reviewer comment (Novelty):
The paper lacks a proper comparison to literature on gradient flow dynamics in linear models. In addition, the citations are often incorrect or correspond to nonexistent papers, as outlined in the section titled Related Work Issues.

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

Reviewer comment (Evaluation):
The comparison of the grokking observed in this model and grokking observed in real networks is still unclear. It is not clear whether the authors think the proposed statistical effect is the phenomenon at play in grokking in real networks.

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

Reviewer comment (Theory):
Many works on grokking report development of specialized weight structures in deep networks temporally coincident with improved generalization error. This form of grokking would correspond to deviation from a linear model, which is outside the explanatory scope of this work. If the authors are not claiming this is a theory of grokking as observed in original settings (like modular arithmetic in transformers), it is unclear in what sense this is an explanation of grokking.

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

Reviewer comment (Theory):
In some grokking experiments with deep networks, the test loss (MSE) increases before later decreasing to its final value (e.g., Davies et al. 2023 Figure 1). It is unclear if this can occur in the linear model where errors in each non-null eigendirection decrease exponentially.

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

Reviewer comment (Evaluation):
The paper could benefit from more discussion on whether the primary aim is to motivate careful future research on grokking. For example, what artifacts due to choice of metric should experimenters be wary of, and how should one distinguish between “improvement in understanding” versus mere statistical noise? Further commentary on this distinction and comparison to prior observed phenomena would improve the paper.

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

Reviewer comment (Novelty):
The authors miss several important works on gradient flow with linear models. Advani & Saxe 2020 derive the training and test error dynamics for the model considered here, and Saxe et al. 2013 derive learning curves for deep linear networks with a fixed dataset. Omitting these citations does not do justice to their contributions.

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

Reviewer comment (Novelty):
Several other papers derive asymptotic performance of linear models trained with Gaussian data in the proportional regime with gradient flow, such as:
   (a) Mignacco et al. 2022 (NeurIPS 2020; JSTAT 2022)
   (b) Paquette et al. 2022

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

Reviewer comment (Experiments):
Figure 5 Bottom row: Is the reason that the dynamics for the tanh networks are well approximated by linear two-layer dynamics because the activation function is linearizable around $\phi(h) \sim h$? What if a different activation function were used, for instance by increasing the “gain” of the tanh nonlinearity like $\tanh(g * h)$ for $g > 1$? For large enough $g$, would we expect the correspondence to the linearized dynamics to break?

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

Reviewer comment (Theory):
Equation 7 is using approximation signs instead of equality. However, if one considered the average case error over random teachers where $T \sim \mathcal N(0,I)$, the equation would be exact. Is there a reason to not focus on averaging over random teachers since the data is isotropic?

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

Reviewer comment (Experiments):
My biggest worry about the paper is that it does not deviate from the toy setting proposed. While the results are good, and the match between analytical results and empirical behavior is close, the entire investigation leaves out realistic applications and the original setting in which the phenomenon was first observed. Furthermore, it seems like the theory does not cover some phenomena associated with Grokking, and it might just be too far from the original setting to yield practical results. The submission would be substantially stronger with more connections to realistic settings or insights that can transfer across, e.g., tasks, datasets, models, etc.

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

Reviewer comment (Novelty):
Delayed generalization in many settings can occur after the generalization loss **increases**, whereas the toy setting presented seems to imply monotonic behavior. Could this mean that the phenomenon studied here is qualitatively different from Grokking in the original setting? There is no fixed and precise definition of Grokking, so this is not a big issue, but I think a distinction might be helpful in understanding feature learning dynamics.

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

Reviewer comment (Theory):
The effect of label noise does not seem to be completely captured by the theory. I have seen scenarios where the label noise is “corrected” at Grokking time, still reaching 100% accuracy. Specifically, one can take a *Grokking on modular addition* setup and assign random labels at a certain rate. Generalization accuracy can still reach 100% eventually. Do you have any idea why that’s the case? Is there a way for the toy model to accommodate such "self-correcting" behavior?

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