Intrinsic Task Symmetry Drives Generalization in Algorithmic Tasks

TL;DR

This paper demonstrates that intrinsic task symmetries are fundamental in driving neural networks from memorization to generalization in algorithmic tasks, revealing a three-stage training process and offering diagnostic tools to predict and accelerate generalization.

Contribution

It introduces the concept that intrinsic task symmetries govern the transition to generalization, supported by a three-stage dynamic and validated across multiple domains.

Findings

Generalization occurs during symmetry acquisition phase.

Representations reorganize into task-aligned geometry.

Symmetry-based diagnostics can predict generalization onset.

Abstract

Grokking, the sudden transition from memorization to generalization, is characterized by the emergence of low-dimensional representations, yet the mechanism underlying this organization remains elusive. We propose that intrinsic task symmetries primarily drive grokking and shape the geometry of the model's representation space. We identify a consistent three-stage training dynamic underlying grokking: (i) memorization, (ii) symmetry acquisition, and (iii) geometric organization. We show that generalization emerges during the symmetry acquisition phase, after which representations reorganize into a structured, task-aligned geometry. We validate this symmetry-driven account across diverse algorithmic domains, including algebraic, structural, and relational reasoning tasks. Building on these findings, we introduce a symmetry-based diagnostic that anticipates the onset of generalization and propose strategies to accelerate it. Together, our results establish intrinsic symmetry as the key factor enabling neural networks to move beyond memorization and achieve robust algorithmic reasoning.