Grokking at the Edge of Linear Separability

We greatly appreciate the reviewer’s positive evaluation and are pleased that they found our perspective on grokking to offer interesting insights. We have thoroughly addressed their concerns below. If they find our responses satisfactory, we hope they will be able to raise their confidence in accepting our work.

Weaknesses

• "The constant label is restrictive”: This appears to be the reviewer’s primary concern, so we will address it thoroughly. We believe this restriction is not as limiting as it may initially seem:

  First, in Section 5, we demonstrate directly that even when the labels are discriminative, grokking can still be observed in the form of delayed generalization - there is a transition from memorization to generalization, albeit imperfect generalization.

  Second, we believe that our model represents a broader class of models that can exhibit grokking for similar underlying reasons. Our focus on the specific case where all labels are the same was primarily motivated by its simplicity and the fact that it presents the phase transition in the cleanest manner.

  To further support this claim, we will briefly describe a related sparse feature classification model that also demonstrates grokking behavior near the critical point with balanced labels: Suppose that the first coordinate of the input, $x_1$, is distributed as a mixture of Gaussians with means $\pm \mu$ and variance $\sigma_1^2$, while all other coordinates $x_i$ for $i>1$ are drawn independently from $\mathcal{N}(0,\sigma^{2})$, as in our original setup. The label of each data point is determined by the sign of the first coordinate. The model is again logistic regression, but with no bias term. By choosing the ratio $\sigma/\sigma_1$ to be sufficiently small, one can observe arbitrarily large grokking time provided we are near the same critical point as our original model, at $\lambda=1/2$. The mechanism behind this behavior is essentially the same as described in the paper: the model initially converges to a memorizing solution (i.e., learning a random separating hyperplane) and only later transitions to the generalizing solution. A numerical demonstration of grokking for this model is included here: https://imgur.com/a/4BuMnls.

  Please let us know if you would like additional analytical or numerical results regarding this model — especially if you believe it could affect your evaluation of the paper. Regardless, we will add to the appendix a brief discussion of this model.

• “Rigor of the results”: We believe that our results are rigorous, note that some complementary proofs are left to the appendices. If any specific issue needs further clarification, please let us know and we will gladly address it.

Questions

• “Why Adam works for $\sigma=1$: We thank the reviewer for this interesting question. First, we will note that numerically one could see the exact same grokking behavior (for $\sigma=1$) not only for Adam, but also for much simpler optimizers such as SignGD (to which Adam behaves similarly, after some time, due to its adaptive nature). Therefore, we need to gain intuition for why $|S|$ grows faster than $b$ in SignGD, leading to grokking. We have $\frac{\partial|S|}{\partial t}=\frac{S}{|S|}\cdot\frac{\partial S}{\partial t}=|\frac{\partial S}{\partial t}|\cos(\alpha)$, where $\alpha$ is the angle between $\frac{\partial S}{\partial t}$ and $S$. Noting that $\frac{\partial S}{\partial t}$ is a vector of $\pm \eta$, we have $|\frac{\partial S}{\partial t}|=\eta\sqrt{d}$, where $d$ is the dimension. In the separable case (or on the verge of being separable), after some time the direction of $S$ saturates so we expect that $\alpha$ would be small. We thus get that $\frac{\partial|S|}{\partial t}\approx\eta\sqrt{d}\gg\eta=\frac{\partial b}{\partial t}$. To sum up, this happens due to the adaptive optimizer nature and the high-dimensionality of the model.
• Linearly separable data with discriminative labels: For $\lambda<1/2$ but using discriminative labels, the model will exhibit late generalization even though the final accuracy differs from 1. For $\lambda>1/2$, similar to the constant-label case, the dynamics will lead us only toward the memorization solution, making the accuracy level stay at values close to its original value of $\approx 1/2$.
• “Which condition is essential for grokking to occur: linear separability or the labels being almost constant?”. In our specific model, being close to the critical point with discriminative labels will result in delayed generalization but not necessarily full generalization. However, more generally, the proximity to the critical point (separability) is the essential point and not the constant labels: the grokking in the balanced-label model that was presented above is a clear demonstration of this fact.

We thank the reviewer for their detailed and valuable feedback.

The reviewer's main concern is the simplicity of the setup and how our results could be generalized to other models (a point also raised by some of the other reviewers). However, the reviewer also noted that “Nonetheless, I believe this is a good first step and should not be used as a basis to reject the paper”. We believe that we have thoroughly addressed the remaining concerns below (please let us know if you think otherwise). In light of this, we hope the reviewer will consider raising their score. We will now address each point in detail.

• Relation to broader literature:

  Novelty with respect to other analytically tractable models - We apologize for not clarifying the difference between our work and previous works. Concretely, the papers cited in the Related works section which pertain to solvable models do not fully analytically solve the dynamics of the models under investigation and do not focus on criticality as a key aspect of the work. Typically, the results are semi-analytical (using DMFT or related techniques), with the only exception being Levi et al. (2023), which we will address in details below. We will revise this section to include this explanation in the final version.

  In contrast to previous works, our model, being simple enough for analytical investigation, allows us to link the fundamental cause of grokking to the existence of critical points. While we cannot yet prove this rigorously, we conjecture that grokking is intimately related to such critical points in settings beyond simple logistic regression. Most other grokking models are more complicated, making the underlying mechanism more elusive. For example, up until recently it was believed that feature learning and weight decay are necessary conditions for grokking, and our model shows this is not the case.

• Weaknesses:

  The simplicity of the model: Rather than attacking the problem from the direction of the ongoing vast research on the canonical examples of grokking (i.e., modular arithmetic), we try to approach it from a different direction by first analyzing its simpler manifestations. As such, we see the simplicity of the model as an advantage rather than a limitation. Of course, the next step must be trying to relate other examples of grokking to the same insights regarding criticality obtained from the simple model. However, we believe that this is a solid starting point, and that the community would benefit from the paper as it stands. We are currently working on using the criticality approach to study other models and also relating it to double descent — these would be addressed separately in our future works.

• Questions:

  1. Differences from Levi et al, 2023: We thank the reviewer for highlighting this important point. We are very familiar with the mentioned paper. The main and most fundamental difference between the two works is that, while Levi et al observe delayed generalization near the critical point, their setup does not display grokking in the regular sense of a transition from a memorizing to a generalizing solution. Simply speaking, their test loss does not exhibit non-monotonicity.

     This limits the scope of their work, since non-monotonicity in the test loss is a hallmark of most known grokking examples. In contrast, our model is the first clean minimal example of grokking that naturally presents grokking both in delayed generalization and a memorization-generalization transition (see for example Fig. 5 in the appendix, where the non-monotonicity is more apparent since the y-axis is not in log-scaled).

     Another obvious difference is that our model deals with cross-entropy loss rather than MSE, and as such, the dynamics and convergence to a solution are vastly different.

     It is, however, interesting to note that their result is also closely tied to criticality: Taking a bit of inspiration from phase transitions in physical systems, we suspect that these two models may belong to different “universality classes” which are classes of transitions that all behave in the same manner in the proximity of the critical point (e.g., have the same critical exponents). The fact that these universality classes are affected only by “general” properties of the model (for example, its symmetries) makes their identification important, as they imply that one could deduce properties of a complex system by studying a much simpler model, as long as it has the same fundamental properties. While further research is certainly needed to verify this, it is part of our motivation for investigating these simpler models.

  2. Typos in the example in Sec. 4: Thank you for pointing these out, both are indeed typos that will be fixed.

• Comments and suggestions:

  We appreciate all three of your suggestions and will clarify these points in the next revision.

We thank the reviewer for their positive and useful feedback, and are glad that the reviewer tends towards accepting our paper.

The reviewer's primary concern is regarding the applicability of our setting to other existing models of Grokking in the literature, which we address below, along with other issues/questions raised.

• “this claim would be far more convincing...”: While we believe our paper stands on its own (see the arguments at the end), we agree that it could benefit from a more detailed reconciliation with existing models. The reviewer has provided some good examples and directions. We will try to address some of them here, while an extension of this discussion will be added as an appendix.

• “..one obvious counter example...”: Our conclusion does not mean that grokking/delayed generalization will be observed only for one ratio of $d/N$, but rather for a range of values near the critical point. However, a divergence of the “grokking time” appears only at the critical point itself. Note that for $N \gg d$ grokking will not exist.

• “Figure 6 of [1] appears...”: We would expect that the grokking time would be smaller for a smaller ratio of $d/N$, but such data does not appear in the figure. Also, Thilak et al. only states that slingshots (not grokking itself) happen - it is not clear if there is a full correspondence. Regardless, the slingshot mechanism itself relies on adaptive optimizers (and $\varepsilon$), so it is not an ideal choice for extending the discussion regarding criticality.

• “being in the 'goldilocks zone..'”. This aligns well with our results. Notice that grokking will not be observed even for $\lambda=1/2$ in our case, unless $\sigma$ is large enough to attract the dynamics toward the memorizing solution before generalizing (see rightmost panel of Fig. 3 in our paper). For more complicated settings, it is likely that other parameters will need to be tuned in order to see the grokking.

• Relation to weight decay: We thank the reviewer for raising this point: It has recently become more apparent that WD is not a necessary condition for grokking, although it can help observe it in certain scenarios, see for example [4]. Our results support this observation: Adding WD will allow us to see grokking to a certain extent even in the non-separable region. We will add a discussion on the effects of WD in our model in the revised manuscript.

• “..the aforementioned work of [2] uses ...” and “Similarly, [5] find...”: Thanks for highlighting these models, which are good candidates to investigate the criticality under broader settings. We will consider addressing it in our upcoming work, but we believe this is beyond the scope of the current work.

• Relation to double descent: Thanks for this important note. We strongly believe such a relation exists, and that our model is a good candidate for investigating it. We are working in this direction and will address it separately in a future work.

• Response to the question: Yes. In fact, these results will hold even for non-symmetric distributions, as this may only change the critical value of $\lambda$.

To sum up, while our paper indeed does not offer a concrete recipe for how these conclusions extend to other models in the literature, we still believe that the community would benefit from its publication for the following reasons:

1. It provides a novel and analytically solvable model that exhibits grokking. Being minimalist, it highlights the true relationship between grokking and criticality — something that is somewhat obscured in more complex models.

2. In contrast to many of the current beliefs, we show that neither a transition from lazy to feature learning, nor WD are necessary conditions for grokking. Instead, the existence of a critical point is the necessary and sufficient requirement. If feature learning is required for this transition, these can coincide.

3. There are strong hints that suggest grokking is related to a criticality [1-3]. Our results support this claim, while offering analytical insights to its mechanism (our approach is somewhat similar to [2], but has much more benefits, see our response to reviewer DzGq for more details).

4. We believe that our results could lay the groundwork for revealing the relation between grokking, criticality, and double descent. In fact, we are currently working on these two directions (the relation to criticality and to double descent), and we intend to address them separately in our upcoming works.

[1] Rubin et al., Grokking as a first order phase transition in two layer networks, 2024.

[2] Levi et al., Grokking in linear estimators – a solvable model that groks without understanding, 2023.

[3] Zhu et al. Investigate the critical data size for language model training through the lens of grokking dynamics, 2024.

[4] Prieto et al. Grokking at the Edge of Numerical Stability, 2025.

We greatly appreciate the reviewer’s thoughtful feedback and positive appraisal of our work. We hope that by addressing the reviewer’s main concerns, they will be able to raise their confidence in accepting our work.

Weaknesses:

1. Limitations of simple models - The reviewer's main concern is the applicability of our setting to more complex models. Our main claim is that our model contains interesting insights regarding the relationship between grokking and criticality. In fact, the value of our model lies in its simplicity and the fact that it is analytically tractable, allowing us to isolate the mechanisms that cause grokking which are significantly harder to disentangle in more complex settings.

   Naturally, the next step will be to extend our results to other models. A characterization of the exact relation between criticality and dramatic phenomena, including both grokking and double descent in complex settings, is currently a work in progress, and will appear in our future work. Nevertheless, we believe the community will benefit from the paper in its present form: please see also the arguments at the end of our response to Reviewer fNB5 (starting from: “To sum up, while our paper...”).

2. Experimental results for non-linear models: We agree that finding nonlinear models where our criticality results still hold would be interesting. However, we believe that extending our criticality analysis to general, nonlinear models warrants a separate work. The goal of this paper is to show a first, solvable model where the key features of grokking are clearly manifest (non-monotonic test loss and delayed generalization), in isolation from other mechanisms, such as weight decay or feature learning. Nonetheless, there are already many works on non-linear models acknowledging that the ratio of samples to dimensions is a crucial factor for grokking. See, for example, Appendix C, where we discuss the relation to canonical examples.

Comments and suggestions:

• Suggestions 1, 2: We thank the reviewer for both of these useful suggestions — we will incorporate these changes into the revised version.

Questions:

1. “Part of the proof is based on the fact that gradient descent dynamics...”: Thank you for bringing up this interesting point. Empirically, we found that in our model, the maximum value that allows generalization is approximately $\eta = 1$, and that below this threshold our gradient-flow calculations fit quite well. Nonetheless, it could be interesting to investigate the effects of large learning rates and how they interact with the grokking phenomenon in this model.

2. “The discussion in Section 5 is good... extend to settings with non-constant labels”: It is indeed possible to extend our results. As is stated in Section 5, for non-constant labels, grokking in the sense of delayed generalization is observed, but not in the sense of perfect generalization. However, to demonstrate that grokking to perfect generalization does not require constant labels, we will present a balanced-label model that exhibits grokking. Suppose the first coordinate $x_1$ is drawn from a Gaussian mixture with means $\pm \mu$ and variance $\sigma_1^2$, while the remaining coordinates $x_{i>1}$ are independently drawn from $\mathcal{N}(0,\sigma^2)$. Labels are determined by $\mathrm{sign}(x_1)$. Choosing a sufficiently small ratio $\sigma/\sigma_1$ yields arbitrarily large grokking times near the critical point $\lambda = 1/2$. Numerical results demonstrating grokking in this model can be found here: https://imgur.com/a/4BuMnls. Note that this also proves that the crucial point for grokking is the proximity to the critical point and not the constant labels, which is a good point raised by Reviewer ZiSB.

Finally, we wish to highlight that our results hold not only for Gaussian inputs but for almost any data distribution (see Appendix H).