# How to Tame Grokking: Representation Geometry as a Control Signal

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 26 references that are in the map.
- First posted on arXiv 13 Jul 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Do Vision Transformers See Like Convolutional Neural Networks?](https://mlsanity.com/paper/2108.08810/) (2021)
- [Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning](https://mlsanity.com/paper/2012.13255/) (2020)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Opening the Black Box of Deep Neural Networks via Information](https://mlsanity.com/paper/1703.00810/) (2017)
- [Measuring the Intrinsic Dimension of Objective Landscapes](https://mlsanity.com/paper/1804.08838/) (2018)
- [Intrinsic dimension of data representations in deep neural networks](https://mlsanity.com/paper/1905.12784/) (2019)
- [Towards Understanding Grokking: An Effective Theory of Representation Learning](https://mlsanity.com/paper/2205.10343/) (2022)
- [Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path](https://mlsanity.com/paper/2106.02073/) (2021)
- [Neural Collapse: A Review on Modelling Principles and Generalization](https://mlsanity.com/paper/2206.04041/) (2022)
- [Grokfast: Accelerated Grokking by Amplifying Slow Gradients](https://mlsanity.com/paper/2405.20233/) (2024)
- [Dimensionality compression and expansion in Deep Neural Networks](https://mlsanity.com/paper/1906.00443/) (2019)
- [Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model](https://mlsanity.com/paper/2504.13292/) (2025)
- [Bridging Lottery Ticket and Grokking: Understanding Grokking from Inner Structure of Networks](https://mlsanity.com/paper/2310.19470/) (2023)

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

- [The Rank the Task Demands: A Causal Rank Law for Matrix Memories Trained on Group Composition](https://mlsanity.com/paper/2609.12259/) (2026)

## Abstract (arXiv)

Grokking is a phenomenon in which neural networks initially memorize training data and only later exhibit strong generalization after prolonged optimization. Despite extensive recent study, the factors influencing the emergence and timing of grokking remain incompletely understood. We investigate the relationship between representation geometry and delayed generalization. We find that dimensionality collapse consistently precedes the onset of grokking in all evaluated settings. Motivated by these observations, we introduce Geometric Dimensionality Regularization (GeomDR), a simple spectral regularizer that modifies the effective dimensionality of hidden representations during training. Across modular addition, modular division, and permutation composition tasks, GeomDR consistently alters grokking dynamics and can substantially accelerate the onset of generalization depending on the intervention schedule and target dimensionality. In several settings, grokking is accelerated by up to 52 times relative to standard AdamW training. Similar qualitative effects are observed in both multilayer perceptrons and transformers. Together, these results suggest that representation geometry can serve as an effective control signal for grokking and provide evidence that geometric interventions offer a practical approach for studying and influencing delayed generalization in neural networks.