# Mechanistic Insights into Grokking from the Embedding Layer

arXiv: https://arxiv.org/abs/2505.15624 · HTML: https://mlsanity.com/paper/2505.15624/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 26 references that are in the map.
- First posted on arXiv 21 May 2025.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Deep Double Descent: Where Bigger Models and More Data Hurt](https://mlsanity.com/paper/1912.02292/) (2019)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Towards Understanding Grokking: An Effective Theory of Representation Learning](https://mlsanity.com/paper/2205.10343/) (2022)
- [Omnigrok: Grokking Beyond Algorithmic Data](https://mlsanity.com/paper/2210.01117/) (2022)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)
- [Explaining grokking through circuit efficiency](https://mlsanity.com/paper/2309.02390/) (2023)
- [Grokking as the Transition from Lazy to Rich Training Dynamics](https://mlsanity.com/paper/2310.06110/) (2023)
- [Grokking modular arithmetic](https://mlsanity.com/paper/2301.02679/) (2023)
- [Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking](https://mlsanity.com/paper/2311.18817/) (2023)
- [Toward Understanding Why Adam Converges Faster Than SGD for Transformers](https://mlsanity.com/paper/2306.00204/) (2023)
- [Deep Networks Always Grok and Here is Why](https://mlsanity.com/paper/2402.15555/) (2024)
- [Unifying Grokking and Double Descent](https://mlsanity.com/paper/2303.06173/) (2023)
- [Grokking at the Edge of Numerical Stability](https://mlsanity.com/paper/2501.04697/) (2025)
- [Grokfast: Accelerated Grokking by Amplifying Slow Gradients](https://mlsanity.com/paper/2405.20233/) (2024)
- [Why Do You Grok? A Theoretical Analysis of Grokking Modular Addition](https://mlsanity.com/paper/2407.12332/) (2024)
- [Weight decay induces low-rank attention layers](https://mlsanity.com/paper/2410.23819/) (2024)
- [Grokking Modular Polynomials](https://mlsanity.com/paper/2406.03495/) (2024)
- [Deep Grokking: Would Deep Neural Networks Generalize Better?](https://mlsanity.com/paper/2405.19454/) (2024)
- [Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model](https://mlsanity.com/paper/2504.13292/) (2025)
- [Deep Learning Through A Telescoping Lens: A Simple Model Provides Empirical Insights On Grokking, Gradient Boosting & Beyond](https://mlsanity.com/paper/2411.00247/) (2024)

## Built on by (most cited first, 6 of 6)

- [Beyond Neural Collapse: Task-Intrinsic Geometry Governs Neural Representations in Modular Arithmetic](https://mlsanity.com/paper/2606.08985/) (2026)
- [Is Grokking Worthwhile? Functional Analysis and Transferability of Generalization Circuits in Transformers](https://mlsanity.com/paper/2601.09049/) (2026)
- [Deciphering Two Training Clocks in Grokking via Deep Linear Network Theory with Conditional ReLU Reduction](https://mlsanity.com/paper/2606.05863/) (2026)
- [At-Grok Is Not Converged:A Measurement-Validity Audit for Grokking Representation Metrics](https://mlsanity.com/paper/2607.06639/) (2026)
- [Tracing the Path to Grokking: Embeddings, Dropout, and Network Activation](https://mlsanity.com/paper/2507.11645/) (2025)
- [Where Grokking Happens: Distributed Utility and Fourier Recoding Without a Module Switch](https://mlsanity.com/paper/2609.17571/) (2026)

## Abstract (arXiv)

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ηEηW propto fracσmax(E)σmax(W) · fracfWfE), 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.