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

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 embeddings, 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 embedding 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.