Title: A Theoretical Framework for Grokking: Interpolation followed by Riemannian Norm Minimisation
Venue: NeurIPS 2025
Accepted: yes
Keywords: grokking; non-convex optimisation; dynamic perturbed systems; theory of neural networks; Riemannian optimisation
Abstract: We study the dynamics of gradient flow with small weight decay on general training losses $F: \mathbb{R}^d \to \mathbb{R}$. Under mild regularity assumptions and assuming convergence of the unregularised gradient flow, we show that the trajectory with weight decay $\lambda$ exhibits a two-phase behaviour as $\lambda \to 0$. During the initial fast phase, the trajectory follows the unregularised gradient flow and converges to a manifold of critical points of $F$. Then, at time of order $1/\lambda$, the trajectory enters a slow drift phase and follows a Riemannian gradient flow minimising the $\ell_2$-norm of the parameters. This purely optimisation-based phenomenon offers a natural explanation for the \textit{grokking} effect observed in deep learning, where the training loss rapidly reaches zero while the test loss plateaus for an extended period before suddenly improving. We argue that this generalisation jump can be attributed to the slow norm reduction induced by weight decay, as explained by our analysis. We validate this mechanism empirically on several synthetic regression tasks.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a theoretical framework explaining how grokking can arise as a result of weight decay. This framework suggests that gradient flow (GF) on loss $F$ with small weight decay has two phases: in the first phase the trajectory quickly converges to a manifold of stationary points of $F$, and in the second phase the dynamics follows the Riemannian GF on this manifold, minimizing the $\\ell_2$-norm of the parameters and slowly converging to an interpolant with a much lower norm, which is believed to generalize better.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\nIn terms of quality:\n* The framework is very general optimization (in particular, no structure is assumed on the parameters). The assumptions seem reasonable for a theoretical treatment. \n* The treatment is (mathematically) rigorous, the claims are made both informally and formally, and proven in the appendix. \n* The theoretical results are supported by simple examples and experiments, covering the standard toy models (quite convincingly).\n* Limitations are discussed very openly. No overclaiming noticed. \n\nIn terms of clarity, the paper is very well-written. The main theoretical results are split into four propositions and explained well. Toy examples are also clearly reported. There is a nicely written review of related work, with connections well-made (rather than just a list of works).\n\nIn terms of significance, this is an important perspective on the reason for grokking, because of its generality. \n\nIn terms of originality, this perspective provides significant improvement over the closest works (such as Lyu et al. 2023, Liu et al. 2022c). \n\n## Weaknesses\n\nIn terms of significance:\n1. This is gradient flow, where weight decay is synonymous with $\\ell_2$ regularization. Transformers are usually trained by smoothed versions of signGD, where its role can be different (for example, AdamW converges to KKT points of $\\ell_\\infty$-norm constrained optimization https://arxiv.org/pdf/2404.04454).\n2. As authors point out, grokking has been observed without weight decay, whereas the explanation in this article is completely based on weight decay.\n3. Relatedly, classification losses are not handled (also a limitation acknowledged by the authors in the conclusion). \n\nAs a minor point, the proof of Proposition 3 in Appendix A.4 got interrupted mid-sentence.\n\nQUESTIONS:\n1. In Figure 3, I could not find the loss curves without weight decay. Did they get lost?\n2. Could the authors clarify a bit the dependence on initialization scale, to provide some connection with previous work that focuses on NTK vs. rich regimes?\n\nLIMITATIONS:\nYes, the limitations are discussed openly.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper gives account of grokking as a dynamical phenomenon that occurs over two timescales: when training with very small weight decay $\\lambda$, gradient flow first (a) converges to very close to whatever solution it was going to anyways, then (b) slowly drifts towards a local optimum of the $\\ell_2$ weight penalty on the manifold of loss minimizers. They prove this for fairly general losses; nothing's specific to deep learning in their analysis.\n\nSTRENGTHS AND WEAKNESSES:\nSome merits of the paper include:\n* it's clearly written and easy to follow.\n* its conclusion is intuitive (this is emphatically a plus, not a minus), and this makes it easy to skip around and get the gist. the various steps link together in an intuitive way.\n* while I didn't check the math, the authors seem to cross their t's and dot their i's; the notation, presentation, etc. seem both mathematically complete and pretty readable, which isn't easy to do.\n\nAll considered, to me, this reads as a cogent and well-presented story, and it's nice to have a formal reference for this intuitive idea. Seems to clear a reasonable bar for quality and usefulness, so I recommend acceptance.\n\n...but ya can't have it all, so here are some critiques, which I share with the aim of improving the paper and explaining why my score isn't higher:\n* This doesn't seem all that conceptually novel. *That's okay* -- too many papers try way too hard to be conceptually novel and force it -- but I'm just noting it. I've heard this qualitative explanation for grokking before, and the story of this paper was basically what I figured was meant: the weight decay kicks in on a longer timescale, so you get slow movement after naive convergence, and this can slide you over to a better solution. It's nice to see this formalized, especially in deep-learning-independent language, but I'm not sure if there supposed to be any surprises or new insights you get from this analysis.\n* I'm not convinced grokking itself is all that interesting or important. I don't count this against the paper for purposes of this review, but it *is* a broader public conversation that should happen somewhere. Why, again, do we think that grokking matters? If anything, the practical utility of this paper seems to be that, having an explanation for grokking, we can stop studying it so intensely, be satisfied that it's just a result of timescale separation, and move on to other things, which we could have done without the formal description of the phenomenon.\n* If there were a mystery of grokking, seems like it'd be in the test performance: why, with unregularized GF, we'd converge to a spurious solution, but then converge to the \"right\" solution with regularization. (My skeptical lean here reflects the fact that this is a much bigger Q and reaches well beyond the phenomenon of grokking.) From this perspective, seems like we've formalized the easy part, but the deeper mystery, except in special cases, is still open.\n\n\nI hope the authors take the fact that my criticisms are fairly high-level to be praise of the paper: I didn't have to spend time critiquing \"implementation details.\"\n\nQUESTIONS:\nI suppose my main questions are conceptual and mostly listed above. The big one (which isn't easy, and which one could and should ask of most papers) is: why does this matter, really? Why does grokking matter, and who does this paper really help? I think a really convincing answer to this actually would get me to increase my score, but that'd be difficult. I'd be curious about the answer regardless of whether or not I'm convinced.\n\nLIMITATIONS:\nyep", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the gradient flow with weight decay in the limit as λ→0. The authors characterize its two-stage behavior: in the first stage it tracks the unregularized gradient flow until it reaches the manifold, and in the second stage it follows the Riemannian flow on that manifold. They use this theoretical framework to explain the phenomenon known as grokking. They verify their findings through synthetic experiments.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**  \n- The paper tackles an interesting problem and provides a new theoretical perspective on the grokking phenomenon.  \n- The writing is fluent, with thorough coverage of related work and background, making it easy to follow.  \n- The simulation experiments convincingly support the authors’ claims.  \n\n**Weaknesses:**  \n1. The main contribution is not prominent. To my knowledge, similar two-stage gradient flow analyses have been widely explored; although the authors’ focus on the weight-decay variant is novel, their theory relies on the λ→0 limit, which substantially weakens the findings. I question the practical relevance of this assumption. Even theoretically, if λ is a nonzero constant, the gradient flow in the first stage may not converge to the manifold at all. Additionally, the proof of the main result (Proposition 2) heavily relies on the definition and properties of the limit flow from prior work, offering limited technical innovation.  \n2. The theory does not fully account for grokking. Under the authors’ framework, the second stage should slowly track the Riemannian flow—implying a plateau and gradual change in validation loss—whereas grokking typically shows a long plateau followed by a sudden drop in validation loss.  \n3. All experiments are synthetic. Experiments on real-world language models are essential.\n\nQUESTIONS:\n1. How do the authors explain the sudden drop in validation loss observed in grokking, rather than the slow Riemannian flow tracking predicted by their theory?  \n2. In Section 4.2, the definition of the limit flow appears to follow Li et al. (2021) rather than being original, yet it is not properly cited.  \n3. Could the authors provide experiments on real-world datasets? Even small-scale ones would be fine.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper lays down purely theoretical justification for grokking phenomenon, a two-phase learning behavior where generalization occurs long after training loss reaches zero. With assumptions on stability, smoothness, and non-pathological behavior of the underlying models, the authors has built a purely theoretical justification of two-phase learning behavior under small l2 weight decay. The paper also supplies small-scale experiments to demonstrate the theoretically discussed behaviors.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n1. The paper presents (almost) model-agnostic explanation of grokking, achieving generalizability and bridges many of the empirically found grokking examples on diverse set of architectures and data (e.g., OmniGrok).\n2. The assumptions underlying the theory is reasonable, and embraces rich set of models.\n3. The paper is well-formatted and pleasurable to read.\n\n## Weaknesses\n\nMajority of my concerns are inquiries rather than criticisms. But here are some points the authors may consider in order to improve the presentation and connection to the practical branch of works.\n\n1. It would be nicer to have a list of real-world functions/models/losses/training scenarios that the underlying assumptions (Assumption 1, 2) hold. It is roughly argued for generality of the assumption throughout page 4. I believe, however, listing concrete examples, e.g., MLP with Lipschitz-continuous activations like smooth rectifying or sigmoidal nonlinearities, sigmoid after quadratics like self-attention layer, etc., will greatly enhance the impact of the theoretical findings, and help the follow-up works to find out where should they start working on.\n2. I believe it will be much better if there is another list that those assumptions do not hold. For example, Line 129-131 and line 328-330 were graceful, but a dedicated list of exclusion maybe in the separate limitation section would be more elegant.\n3. The first report on grokking phenomenon was on Transformer network, which contains not only residual MLP but also residual multi-head attention. So it will be nicer to have this application explained by the proposed theory, too.\n\nOverall, I believe the paper is already in a good shape regarding the theoretical novelty. However, I have several concerns related to presentation and demonstration. I will raise my score if those issues are resolved. Please also refer to the questions section.\n\nQUESTIONS:\nThese are the questions that I do not think to be appropriate to be classified as the weaknesses but do still believe to be important. Below are request for clarifications of my possible misunderstandings.\n\n1. I may have lost but is t≈1/λ transition repeatedly appearing throughout the manuscript justified? Is this transition boundary sharp or gradual in practice?\n2. How small the weight decay should be in order to be applied by this theory? Do the authors have a clue when does the grokking start to not happen as we increase the weight decay OR decrease the initial weights?\n\nThe next question is from my pure curiosity and therefore the answer will not affecting the overall scores.\n\n1. In the Line 347-349 and the referred Appendices C and D.2, the authors have argued that grokking may promote sparse estimators. From my point of view, grokking seems to “correct” the bad initialization at the beginning of the training. As demonstrated in OmniGrok, sufficiently small weight initializations can quickly set the model to the so-called “rich region”, and **de-grok** the model. This implies that we have to be careful about the features that seem to be **unlocked** by the grokking phenomenon. Is this the “grokking” phenomenon that “unlocks” certain features of the model that would not be possible without “large weight initialization & grokking”, or is it just the problem-intrinsic feature that can also be achieved by **de-grokked** models?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-03"}