Intrinsic Task Symmetry Drives Generalization in Algorithmic Tasks

Grokking, a sudden transition from memorization to generalization, has been closely linked to the emergence of low-dimensional representations; yet the mechanism driving this organization remains elusive. Here, we propose that intrinsic task symmetries are the key drivers of grokking, inducing structured geometries in representation space. Our analysis reveals a consistent three-stage training dynamic: (i) data memorization, (ii) intrinsic symmetry acquisition, and (iii) geometric organization. We show that generalization emerges during the symmetry acquisition phase, and subsequently the embedding space organizes into a low-dimensional structured geometry. We validate this mechanism across diverse algorithmic domains, spanning algebraic (modular arithmetic), structural (graph metric completion), and relational (comparison) reasoning tasks. Leveraging these insights, we formulate a symmetry-based criterion for generalization and propose symmetry- and geometry-prompting training strategies that can accelerate generalization. Together, our results establish intrinsic symmetry as a central mechanism enabling neural networks to move beyond memorization and achieve robust algorithmic reasoning.

Deep learning models sometimes exhibit a surprising behavior: after a long period of simply memorizing training examples, they suddenly begin to generalize and solve new problems correctly. Although this phenomenon, known as grokking, has been widely observed, its underlying cause has remained unclear. In this work, we show that the key factor behind this transition is the discovery of hidden symmetries within a task. As training progresses, neural networks first memorize examples, then learn these underlying symmetries, and finally organize their internal representations into simple geometric structures. Importantly, we find that the ability to generalize emerges when the network acquires the task's symmetries, before these geometric structures are fully formed. We demonstrate this behavior across several reasoning tasks, including arithmetic, graph-based problems, and comparison tasks. Based on these findings, we propose practical training methods that encourage networks to learn symmetries and structured representations more efficiently, leading to faster and more reliable generalization. Our results provide a new explanation for how neural networks move beyond memorization and develop robust reasoning abilities.