Title: Mechanistic Insights into Grokking from the Embedding Layer
Venue: NeurIPS 2025
Accepted: no
Keywords: Embedding learning

Token frequencey

Coupled system
Abstract: Grokking, a delayed generalization in neural networks after perfect training performance, has been observed in Transformers and MLPs, but the components driving it remain underexplored. We show that embeddings are central to grokking: introducing them into MLPs induces delayed generalization in modular arithmetic tasks, whereas MLPs without embeddings can generalize immediately. Our analysis identifies two key mechanisms: (1) Embedding update dynamics, where rare tokens stagnate due to sparse gradient updates and weight decay, and (2) Bilinear coupling, where the interaction between embeddings and downstream weights introduces saddle points and increases sensitivity to initialization.  
To confirm these mechanisms, we investigate frequency-aware sampling, which balances token updates by minimizing gradient variance, and embedding-specific learning rates, derived from the asymmetric curvature of the bilinear loss landscape. We prove that an adaptive learning rate ratio, \(\frac{\eta_E}{\eta_W} \propto \frac{\sigma_{\max}(E)}{\sigma_{\max}(W)} \cdot \frac{f_W}{f_E}\), mitigates bilinear coupling effects, accelerating convergence. Our methods not only improve grokking dynamics but also extend to broader challenges in Transformer optimization, where bilinear interactions hinder efficient training.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the underexplored role of embedding layers in the phenomenon of grokking—the delayed generalization in neural networks after perfect training performance. The authors focus on MLPs with and without embedding layers trained on modular arithmetic tasks and identify two core mechanisms:\n1. Sparse embedding updates due to infrequent token sampling, leading to stagnation.\n2. Bilinear coupling between embeddings and first-layer weights, which complicates optimization.\n3. They propose two practical solutions: Frequency-aware sampling to ensure balanced token updates; Embedding-specific learning rate scaling (Adam-LR) based on singular value and frequency analysis.\nExperiments on synthetic algorithmic datasets confirm their hypotheses and show accelerated grokking and better generalization using their methods.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThis paper provides a compelling and novel perspective on the phenomenon of grokking by focusing on the underexplored role of embedding layers in MLPs. Through both theoretical analysis and carefully designed experiments on modular arithmetic tasks, the authors identify two key mechanisms—sparse embedding updates and bilinear coupling with downstream weights—that delay generalization. These insights are well-grounded in optimization theory, including gradient dynamics and curvature analysis. The proposed solutions, namely frequency-aware sampling and the Adam-LR optimizer with embedding-specific learning rate scaling, are simple yet effective. The empirical results show that these methods accelerate convergence and improve generalization, demonstrating clear practical benefits. Moreover, the paper connects its findings to broader challenges in Transformer optimization, suggesting relevance beyond the narrow setting of algorithmic tasks.\n\nWeaknesses:\nThe paper is limited in scope due to its exclusive focus on synthetic datasets and MLP architectures. While the modular arithmetic tasks offer a controlled environment for analysis, it remains unclear how well the proposed methods generalize to real-world settings or more complex architectures like Transformers. The empirical validation would benefit from benchmarks on natural language or vision tasks where embedding layers are critical. Additionally, although the paper touches on the impact of weight decay, learning rate, and sampling strategies, more detailed ablation studies could help disentangle their individual effects. Finally, while the learning rate scaling heuristic is well-motivated, the absence of a principled criterion for selecting the optimal scaling factor may limit its adoption in practice.\n\nQUESTIONS:\n1. Have you tested your embedding-specific learning rate scaling in standard NLP or vision settings?\n2. Would explicit regularization of the embedding matrix (to control rank) be as effective as your adaptive optimizer?\n3. How sensitive are your conclusions to the assumption of fixed token probabilities? \n4. Your analysis focuses on a shallow MLP with only two layers. Have you tried validating your hypotheses (e.g., embedding-gradient imbalance and bilinear coupling) on deeper MLPs or other architectures\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
The paper investigates the role of embedding layers in grokking in multi-layer perceptron networks (MLPs) for modular arithmetic tasks. The work highlights that embedding layers and their learning dynamics play a central role in grokking behavior. Rare tokens receive fewer gradient updates, leading to slower convergence of their embeddings, which partly explains why generalization is delayed. In addition, embeddings also interact with the first layer weights, and this interaction further impacts embedding layer gradients. To address these, the paper proposes two strategies to accelerate grokking, refined sampling and using embedding specific learning rates via a custom optimizer (Adam-LR). In the limited MLP based setup, these techniques significantly speed up grokking across modular arithmetic tasks.

STRENGTHS AND WEAKNESSES:
Strengths:

1. The paper provides a concrete mechanistic explanation for grokking by looking at the interaction between embeddings and downstream weights. 
2. The propose modification to the learning rate for embeddings is shown to accelerate grokking, which is good.
3. The experimental design allows for clear analysis of the phenomenon.

Weaknesses:
1. The main weakness is the limited scope of the study - the experiments are confined to MLPs. There is no empirical evidence that the findings would generalize to more complex architectures and attention blocks.
2. The analysis is done on small synthetic tasks which does not capture the complexity of real world data.

QUESTIONS:
How does the embedding-gradient stagnation compare to stagnation in other parts of the network? Could similar delayed dynamics occur elsewhere?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper hypothesizes that embedding layers are at the heart of grokking: unlike MLPs without embeddings (which can generalize immediately on simple tasks), embedding-equipped models exhibit pronounced delayed generalization, Two issues are identified: (i) rare tokens receive gradient updates only when sampled in a batch and otherwise, their embeddings decay with weights, causing delay in convergence; (ii) the bilinear interaction between the embedding matrix $E$ and the first-layer weights $W$ creates saddle points and increases sensitivity to initialization. The paper then adopts frequency-aware sampling and scales the embedding LR to improve grokking behaviour. A new optimizer, Adam-LR also addressed the imbalance in updates, Authors also uncover emergent basis functions, aligned with the literature.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- The paper identifies an interesting point regarding embedding layers / networks and their relations to grokking.\n\n### Weaknesses\n\n- I think the introduction does not do a good job in presenting the intuition as to why *embedding layers* should be the central objects to look into. Is this a coincidence / chance or is it somehow informed? The intro should also do a better job of separating MLPs and transformers. If we are interested in grokking in both architectures, this should be clear. \n\n- For the contributions: wouldn't the dynamic updates and the coupling also impact the training accuracy? Why should they act differently on train and test?\n\n- The paper focuses on transformers which are greatly used in practice. However, the data / task settings are still toy. \n\n- In Figure 2, I disagree that generalization doesn't occur for multiplication. How does the paper exactly define generalization? I would define it as the gap between test and training accuracies. As such, both plots seem to generalize fairly well: the gap is small; the test accuracy matches the train accuracy. To me, generalization doesn't say anything about the solution quality (training / testing accuracy). It is about closing the gap. Maybe this could be clarified. \n\n- The paper seems to formulate everything continuously whereas the token space can have a discrete / quantized structure. Can the paper say something about this? Do the techniques still apply? \n\n- The embedding dynamics follow a gradient-descent like update, whereas the paper largely talks about and emphasizes ADAM optimizer. In fact a new optimizer in the form of an intervention, called ADAM-LR, is developed. I think there is an unaddressed modeling gap here. How do we ensure that the observations as well as theories generalize?   \n\n- I find it hard to create a coherent story out this paper since the data-splitting strategies do not really fit into the proposed framework of embedding layers. How do these fronts interact and create a whole? These sampling strategies are not even combined and evaluated with the proposed optimizer.\n\n- The singular value analysis is interesting and there is a rich literature on this. One interesting paper would be from Yunis et al.: \n\nYunis et al. Grokking, Rank Minimization and Generalization in Deep Learning. ICML 2024 Workshop on Mechanistic Interpretability,\n\nI think the paper could do a better job exposing and grounding the current work in conjunction with that literature. Maybe some discussion could involve flat minima theories, e.g.:\n\nMulayoff & Michaeli. Unique Properties of Flat Minima in Deep Networks. ICML 2020.\n\nQUESTIONS:\n- How can this work be connected to the numerical issues highlighted in Prieto et al.? Is grokking a feature or a bug here? Can the optimizer instabilities be explained/predicted by the numerical problems arising during training?\n\n- The paper focuses on transformers which are greatly used in practice. However, the data / task settings are still toy. Could the authors comment on this gap? Are we ready to characterize grokking in real scenarios? \n\n- In Fig. 3, why are the plots not ordered by steps? Is the order particular? What should we observe there?\n\n- What about other eigenvalues and not just the maximum? How does the behaviour depend on that?\n\n- How strong is the assumption that the gradient stabilizes to a constant value? For a relatively empirical paper, I find this to be quite strong. This is a fundamental difference between GD and SGD.\n\n- Could the paper better discuss the large batch size requirement? It seems that it is almost framed as a precondition - however most optimizers like ADAM do not require this by default.\n\nPlease also address the other questions in the weaknesses section.\n\nLIMITATIONS:\nThe limitations are mentioned in the discussion section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**Grokking** refers to the surprising phenomenon where neural networks, after reaching perfect training accuracy, continue to reduce training loss without improving test performance, until they suddenly transition to strong generalization. Many prior works have attempted to explain grokking in simplified models by invoking high-level ideas such as feature learning, phase transitions in representation geometry, or emergent circuits. However, these explanations often remain abstract and disconnected from specific architectural components or optimization dynamics.\n\nThis paper presents a novel and grounded perspective on understanding by focusing on the role of **embedding layers in MLPs** trained on algorithmic tasks, such as modular arithmetic. The authors show that embedding layers introduce structural dynamics that delay generalization. Specifically, MLPs without embeddings generalize immediately on modular addition, while the same task exhibits grokking behavior when embeddings are present. For modular multiplication, generalization requires embeddings, highlighting their expressiveness and structural role.\n\nThe paper identifies two concrete mechanisms that contribute to this effect:\n\n1. **Sparse embedding updates**: Rare tokens appear infrequently in training batches, so their embeddings receive few gradient updates and mostly decay due to weight regularization. As a result, those embeddings converge slowly. The authors support this explanation by varying the sampling strategy, showing that more uniform sampling leads to earlier generalization.\n\n2. **Bilinear coupling**: The embedding matrix and the first-layer weights interact bilinearly, creating a more complex optimization landscape with saddle points and higher sensitivity to initialization, especially in the scale of the embedding weights.\n\nTo mitigate these effects, the authors propose two interventions:\n- **Frequency-aware sampling** to equalize the update rates across tokens.\n- An **embedding-specific learning rate**, derived from the relative singular values and update frequencies of the embedding and weight matrices, to counteract bilinear coupling.\n\nThey implement this strategy in a modified optimizer, **Adam-LR**, which accelerates grokking and stabilizes training across several modular tasks. The authors link these improvements to more balanced Hessian spectra between embeddings and downstream weights.\n\nWhile the experiments focus on MLPs, the insights may extend to Transformer models, which also involve bilinear interactions through attention mechanisms. The paper identifies this as an important direction for future work.\n\nSTRENGTHS AND WEAKNESSES:\nThis paper provides a welcome shift from abstract explanations of grokking to a more grounded, mechanistic account that centers on the training dynamics of the embedding layer. While previous work has focused on simplified models and high-level explanations—such as phase transitions, feature learning, or emergent circuits—this paper stands out for its attention to low-level architectural and optimization features that are present in widely used neural networks.\n\nThe methodology is sound, the experiments support the hypotheses, and the proposed optimizer is both simple and effective. The analysis of curvature, initialization sensitivity, and bilinear interactions offers insights that are likely relevant beyond MLPs, particularly for Transformer architectures where similar optimization challenges arise.\n\n**Overall**, this is a strong submission that contributes novel insights into the grokking phenomenon by linking it to concrete architectural components and optimization behaviors. The work is likely to interest both theoretical and applied communities in deep learning.\n\nOne limitation of the submission, which could be addressed with additional space or revisions during the camera-ready phase, is its limited discussion of the existing explanations about grokking. A more detailed comparative discussion—possibly in the style of a brief survey—could help readers understand how the mechanisms proposed here relate to or complement prior theories developed in both neural and non-neural settings. A discussion in the non-neural setting might also highlight how much of grokking can be attributed to the embeddings (cf. [this paper](https://arxiv.org/abs/2407.20199)). \n\nAnother limitation is that I would have liked to see experiments on more tasks and architectures. Still, I understand that the insights (as well as the correspondingly more reasonable optimization algorithm) might change for a different task.\n\nQUESTIONS:\n1. (**Critical to me raising my score**) Can the authors include in their response a more detailed and technical discussion about the existing explanations about grokking, and how they relate to the insights presented in this paper? \n\n2. (**Not critical but useful**) If the authors ran other experiments, but did not include them in the submission, it would be good to know about those during the rebuttal phase.\n\n2. (**Just a curiosity**) Does looking at the evolution of the entire spectrum as opposed to just the top singular value offer more insight into the experiment of Figure 7? Specifically, I am curious if the authors have reviewed the observations about weight decay and grokking established in [this paper](https://arxiv.org/abs/2408.11804).\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies grokking behaviour, with a particular focus on embedding layers. It proposes that the embedding layer should be optimized with a larger learning rate than the other layer in a two-layer MLP, based on an analysis of gradient dynamics.\n\nSTRENGTHS AND WEAKNESSES:\n- It's not very surprising that in the problem you study, your attempt at optimal per-token learning rates didn't do much; there isn't much gap in how often the tokens appear, since at the distribution level they're all uniform. This issue of less-frequent updates for rare tokens is far more problematic in language settings; for example, [Li et al. (ICLR 2022)](https://openreview.net/forum?id=ibqTBNfJmi) study it quite directly.\n\n- [Hayou et al. (ICML 2024)](https://arxiv.org/abs/2402.12354) do a very similar, but more thorough, analysis of LoRA – whose important structure here is roughly the same – and propose a similar learning rate adaptation technique. The followup [Hayou et al. (arXiv 2025)](https://arxiv.org/abs/2506.15025) – fully acknowledging that this work was apparently parallel to yours and went online only after this deadline – does a similar analysis specifically for embedding layers.\n\n- The point here about the importance of initialization scale has also been observed by prior work on grokking; for instance, [15] relate it to the speed at which optimization departs the \"kernel regime.\" It would be good to understand whether the patterns you identify here are the particular way in which that effect applies in this setting, or if they are counter to these prior explanations.\n\n- Given that you are fundamentally concerned about the bilinear structure of an embedding layer followed by weights, it's surprising to me that you didn't mention (let alone pull anything from) the extensive literature on deep linear models. One reasonable entry point might be [Dominé et al. 2024](https://arxiv.org/abs/2409.14623). The analysis of [Telgarsky and Son](https://openreview.net/forum?id=lYQLwP9c9S) also feels like it is probably relevant.\n\nOverall: grokking on modular arithmetic is, and always has been, a toy problem. This paper adds some amount of explanation to two-layer MLPs on this problem, but not much, and no real attempt is made to position it with respect to the last several years of understanding grokking. To the extent that the insights here generalize, they probably generalize to embedding layers of transformers in language models; the present work does not address whether these insights should generalize to those settings at all, while they have been and continue to be studied in those cases. I think there is something interesting to learn from the bilinear point of view there – perhaps even things not simultaneously found by Hayou et al. (arXiv 2025) – but this paper doesn't address them.\n\n## Smaller points\n\n- The discussion at the bottom of page 3 about how \"there is no statistical 'distribution' in the conventional sense\" is mostly misleading. There absolutely is a distribution, and this is entirely in the regime of e.g. traditional learning theory; it just happens to be a finitely-supported distribution. The discussion here seems to imply that for instance, one might train the model on all possible input-output pairs, but this is not what leads to grokking; rather, grokking occurs within a range of how many of the possible training points you observe.\n\n- The discussion at the top of page 4 about the isomorphism between modular addition and multiplication was present in the original grokking paper [19] (second paragraph of their Section 3.2).\n\n- The analysis in (5) and (6) assumes that the loss is $\\beta$-smooth, but you specifically call out just below (1) that you are studying the model only with ReLU activations, which are _not_ $\\beta$-smooth.\n  - The result you show in (5) and (6), which is an immediate corollary of the very well-known \"descent lemma\", does have a known ReLU network analogue; see for example [this result](https://mjt.cs.illinois.edu/dlt/#smoothness-inequality-adapted-to-relu).\n\n- The limit where the updates \"behave like a continuous system\" is known as \"gradient flow\" in the literature.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nSocietal impact is not really applicable. The limitations of the insights given by the paper are not thoroughly addressed, and that makes up the bulk of my complaints above.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-10"}