The Geometric Origin of Grokking: Accelerating Generalization via Active Structural Reorganization

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.

When AI models train on data, they sometimes seem to "get it" all at once — suddenly solving problems they couldn't crack before. This phenomenon, known as grokking, has puzzled researchers because we didn't know what was actually happening inside the model. We looked inside the model and found that generalization is triggered by a reorganization of how the model internally encodes information. Think of it like a messy desk suddenly sorting itself into neat piles: before the reorganization, the model treats every input as an unrelated symbol; afterward, it has figured out the underlying structure connecting them. Building on this insight, we designed a lightweight training add-on called R2G Loss that actively nudges the model toward this organized state — like giving the messy desk a push. It requires no extra data and no changes to the model itself. Across a range of math and language tasks, R2G dramatically increases the chance that a model successfully generalizes, even rescuing runs that would otherwise never learn at all.