The Geometric Origin of Grokking: Accelerating Generalization via Active Structural Reorganization
Kefei Tao, Zhang Zhang, Mingze Qi, Xiaojun Duan
Abstract
Grokking, the phenomenon where models suddenly generalize long after overfitting training data, remains a puzzling challenge in neural network dynamics. Through mechanistic analysis, we find that this transition is fundamentally driven by a structural reorganization of token representations, with the onset of grokking entailing a shift toward a well-defined geometry, and reveal the model’s distinct understanding of data’s dual characteristics. Building on these geometric insights, we propose R2G (Repel-to-Grokking) Loss, an active intervention that reshapes the representation manifold by enforcing structural repulsion. The versatility of R2G is empirically validated in both algorithmic and linguistic tasks, while our theoretical analysis and ablation studies jointly demonstrate that angular reorganization is the primary driver of grokking. Our work offers a novel mechanistic perspective on the evolution of grokking and provides a useful tool for enhancing model efficiency and reliability.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext ae85b198-257e-45e7-95a7-9c2bcd770592Builds on7
- Towards Understanding Grokking: An Effective Theory of Representation LearningZiming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud et al.NeurIPS 2022 · 299 citations
- The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural NetworksZiqian Zhong, Ziming Liu, Max Tegmark, Jacob AndreasNeurIPS 2023 · 181 citations
- A Toy Model of Universality: Reverse Engineering how Networks Learn Group OperationsBilal Chughtai, Lawrence Chan, Neel NandaICML 2023 · 144 citations
- Progress measures for grokking via mechanistic interpretabilityNeel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith et al.ICLR 2023 · 54 citations
- Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasksTianyu He, Darshil Doshi, Aritra Das, Andrey GromovNeurIPS 2024 · 52 citations
Related papers
- Omnigrok: Grokking Beyond Algorithmic DataZiming Liu, Eric J. Michaud, Max TegmarkICLR 2023 · 8 citations
- Is Grokking a Computational Glass Relaxation?Xiaotian Zhang, Yue Shang, Entao Yang, Ge ZhangNeurIPS 2025 · 8 citations
- Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via GrokkingTing Han, Linara Adilova, Henning Petzka, Jens Kleesiek et al.NeurIPS 2025 · 9 citations
- Explaining Grokking and Information Bottleneck through Neural Collapse EmergenceKeitaro Sakamoto, Issei SatoICLR 2026 · 5 citations
- Intrinsic Task Symmetry Drives Generalization in Algorithmic TasksHyeonbin Hwang, Yeachan ParkICML 2026 · 1 citation
