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.
Figures & tables
Figure 1: Modular Arithmetics Embeddings
Figure 2: Graph Metric Completion and Comparison Embeddings.
Figure 3: Type of Graphs.
Figure 4: Full PCA visualization of modular addition embeddings.
Figure 5: Symmetric patterns within a dataset lead the network to learn intrinsic symmetry.
Figure 6: Three-stage dynamic of grokking for modular addition task. (1) Memorization : The model first memorizes as it quickly reaches perfect train accuracy. (2) Symmetry Acquisition Then with the symmetry criterion decreasing to certain threshold, perfect test accuracy is reached. (3) Geometric Organization Finally, as the representation collapses into lower-dimension, embedding forms a geometric organization.
Figure 7: Symmetry violation metric vs. test accuracy across different seeds. A threshold in symmetry violation separates Solved from Unsolved runs, supporting symmetry alignment as a predictor of generalization.
Figure 8: We train a two-layer MLP on modular addition using an 80% training fraction, without weight decay.
Figure 9: Comparison of convergence speed on Modular Addition, Graph Completion (Lattice), and Comparsion (2D) task with Geometry-promoting priors ( nuclear , entropy , Lipschitz ); Symmetry-prompting losses ( commutativity / associativity /etc); Optimizer-based accelerants ( GrokFast ). Markers denote the first point where test accuracy reaches 100%. Consistency refers to using both Commutativity and Associativity.
Appendix figures & tables46 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Grokking Dynamics in Modular Addition.
Figure 11: Grokking Dynamics in Modular Subtraction.
Figure 12: Grokking Dynamics in Modular Squared Sum
Figure 13: Grokking Dynamics in Modular Cubed Sum
Figure 14: Grokking Dynamics in Modular Multiplication
Figure 15: Grokking Dynamics in Modular Affine–bilinear
Figure 16: Grokking Dynamics in Graph Metric Completion: Cycle
Figure 17: Grokking Dynamics in Graph Metric Completion: Path
Figure 18: Grokking Dynamics in Graph Metric Completion: Cylinder
Figure 19: Grokking Dynamics in Graph Metric Completion: Hypercube
Figure 20: Grokking Dynamics in Graph Metric Completion: Lattice
Figure 21: Grokking Dynamics in Graph Metric Completion: 3D Lattice
Figure 22: Grokking Dynamics in 2D Comparison
Figure 23: Grokking Dynamics in 3D Comparison
Figure 24: Maximum Symmetry Violation and Generalization Status for modular arithmetic tasks.
Figure 25: Symmetry Loss vs. Test Accuracy for Graph Metric Completion Tasks
Figure 26: Symmetry Loss vs. Test Accuracy for Comparison Tasks
Figure 27: PCA projection of embeddings of model trained on modular addition task.
Figure 28: PCA projection of embeddings of model trained on modular subtraction task.
Figure 29: PCA projection of embeddings of model trained on Task 3.
Figure 30: PCA projection of embeddings of model trained on Task 4.
Figure 31: PCA projection of embeddings of model trained on modular multiplication task.
Figure 32: PCA projection of embeddings of model trained on Task 6.
Figure 33: PCA projection of embeddings of model trained on the Path graph metric completion.
Figure 34: PCA projection of embeddings of model trained on the Cycle graph metric completion.
Figure 35: PCA projection of embeddings of model trained on the Cylinder graph metric completion.
Figure 36: PCA projection of embeddings of model trained on the Hypercube graph metric completion.
Figure 37: PCA projection of embeddings of model trained on the 2D Lattice graph metric completion.
Figure 38: PCA projection of embeddings of model trained on the 3D Lattice graph metric completion.
Figure 39: PCA projection of embeddings of model trained on the 2D Comparison .
Figure 40: PCA projection of embeddings of model trained on the 3D Comparison .
Figure 41: Convergence speed on modular addition when using various regularization.
Figure 42: Convergence speed on modular subtraction when using various regularization.
Figure 43: Convergence speed on modular squared sum when using square sum.
Figure 44: Convergence speed on modular cubed sum when using various regularization.
Figure 45: Convergence speed on modular multiplication when using various regularization.
Figure 46: Convergence speed on modular affine–bilinear when using various regularization.
Figure 47: Convergence speed on cycle graph metric completion when using various regularization.
Figure 48: Convergence speed on path graph metric completion when using various regularization.
Figure 49: Convergence speed on cylinder graph metric completion when using various regularization.
Figure 50: Convergence speed on hypercube graph metric completion when using various regularization.
Figure 51: Convergence speed on 2d lattice graph metric completion when using various regularization.
Figure 52: Convergence speed on 3d lattice graph metric completion when using various regularization.
Figure 53: Convergence speed on Comparison (2D) when using various regularization.
Figure 54: Convergence speed on Comparison (3D) when using various regularization.
Figure 55: Comparison of embedding initialization and output scaling on Modular Addition, 1-Layer Transformer
Grokking occurs when a model achieves high training accuracy but generalization to unseen test points happens long after that. This phenomenon was initially observed on a class of algebraic problems, such as learning modular arithmetic (Power et al., 2022). We study grokking on algebraic tasks in a class of feature learning kernels via the Recursive Feature Machine (RFM) algorithm (Radhakrishnan et al., 2024), which iteratively updates feature matrices through the Average Gradient Outer Product (AGOP) of an estimator in order to learn task-relevant features. Our main experimental finding is that generalization occurs only when a certain symmetry in the training set is broken. Furthermore, we empirically show that RFM generalizes by recovering the underlying invariance group action inherent in the data. We find that the learned feature matrices encode specific elements of the invariance group, explaining the dependence of generalization on symmetry.
Marcel Tomàs Bernal, Neil Rohit Mallinar, Mikhail Belkin
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.
Why do neural networks memorize algorithmic training data long before they generalize? We present a geometric case study demonstrating that, on tasks where generalization requires discovering structured low-dimensional circuits, the memorization-generalization delay is driven by radial inflation of hidden representations under cross-entropy optimization. We formalize a radial-angular decomposition of activation-space dynamics and derive three testable propositions: (i) that penalizing radial inflation induces anisotropic, data-dependent weight regularization; (ii) that it suppresses radial gradient energy below the isotropic random baseline, forcing predominantly angular updates; and (iii) that it biases convergence toward flatter minima. To empirically validate these propositions, we study a single-hyperparameter norm penalty that softly constrains activations to a sqrt(d)-radius hypersphere. On modular arithmetic, this penalty accelerates grokking up to 6x across MLPs and Transformers, and halves training steps for a 10M-parameter nanoGPT on 3-digit addition.