cs.AIMar 31, 2026
SaveGrokking From Abstraction to Intelligence
Abstract
Grokking in modular arithmetic has established itself as the quintessential fruit fly experiment, serving as a critical domain for investigating the mechanistic origins of model generalization. Despite its significance, existing research remains narrowly focused on specific local circuits or optimization tuning, largely overlooking the global structural evolution that fundamentally drives this phenomenon. We propose that grokking originates from a spontaneous simplification of internal model structures governed by the principle of parsimony. We integrate causal, spectral, and algorithmic complexity measures alongside Singular Learning Theory to reveal that the transition from memorization to generalization corresponds to the physical collapse of redundant manifolds and deep information compression, offering a novel perspective for understanding the mechanisms of model overfitting and generalization.
Figures & tables
Figure 1: Grokking on modular arithmetic task mod 97.
Figure 2: CMA for architectural evolution during emergence and ablation experiment for attention heads. Sub-figures (a/b/c/d) represent modular arithmetic tasks with .
Figure 3: The emergence of generalization is mechanistically characterized by (a) the transition to a group-theoretic structure and (b) the spectral localization of weight energy in the Fourier domain.
Figure 4: The curve tracks the global BDM complexity of the model parameters throughout training. During the emergence phase, the algorithmic complexity of the model weights drops sharply, and the internal parameter structure becomes noticeably block‑structured, with different colors indicating discretized parameters after quantization.
Figure 5: Left: Generalization (green) aligns with a collapse of geometric-complexity proxies (red). Right: Spectral weights transition from high-entropy noise to sparse, low-dimensional structured patterns via the Occam Gate (diagonal-like structure is most directly expected for addition/subtraction; other operations may require reindexing such as discrete-log ordering).
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
| Item | Specification |
|---|---|
| Task | Modular arithmetic over ; operations ; for exclude . |
| Input/target format | Input tokens: (length 4); target token . |
| Vocabulary | , . |
| Dataset construction | Enumerate all valid input pairs. Total pairs: for and for . |
| Train/test split | Random split by input pairs with a fixed PRNG seed. Default split: 50%/50%. For : 4704 train / 4705 test; for : 4656 train / 4656 test. Split seed = 0. |
| Evaluation | Report train and test accuracy (exact match) and cross-entropy loss. Evaluate every 100 steps and additionally at checkpoints 0.1k/1k/10k/100k steps. |
Table 1: Dataset and evaluation protocol for all experiments.
| Item | Specification |
|---|---|
| Architecture family | Decoder-only Transformer (GPT-2 style), 48 layers. |
| Context length | (inputs are length-4 token sequences). |
| Model width | Embedding/hidden size ; MLP size (MLP ratio 4). |
| Attention | Num heads ; head dim . |
| Dropout | Attention/MLP dropout = 0.0. |
| Normalization | Pre-LN with LayerNorm. |
Table 2: Model architecture details (fixed across runs unless explicitly stated).
| Item | Default training hyperparameters |
|---|---|
| Optimizer | AdamW. |
| AdamW params | . |
| Learning rate | Peak LR ; min LR . |
| LR schedule | Linear warmup for 1000 steps (1% of total), then cosine decay to . |
| Batch size | Global batch size = 512 sequences; gradient accumulation = 1. |
| Weight decay | for the main experiments; applied as decoupled weight decay (AdamW). |
Table 3: Default optimization setup for the 48-layer Transformer.
| Analysis / ablation | What is varied and what is reported |
|---|---|
| Seed sensitivity | For each key figure, repeat over the seed set in Table 3 ; plot mean trajectory and shaded std (train/test). |
| Weight decay sweep | ; report grokking onset distribution across seeds and the fraction of runs that grok. |
| Learning-rate sweep | (schedule unchanged); report whether grokking occurs and time-to-grok. |
| Warmup sweep | Warmup steps (total steps fixed at 100k); report impact on memorization length and generalization onset. |
| Batch-size sweep | Global batch size ; keep total steps fixed; report variance and onset shift. |
| Early-stopping check | Not used in main experiments; when enabled for a control, monitor test accuracy with patience 10,000 steps. |
Table 4: Sensitivity analyses and ablations reported for robustness.
Figure 6: Benchmarking BDM against baselines and controls. Left (B.1): Comparison with Standard Proxies. BDM (red) aligns closely with Block Entropy, Spectral Entropy, and Effective Rank, confirming it captures genuine complexity reduction. BDM exhibits a slightly sharper transition, indicative of its sensitivity to local geometric patterns. Right (B.2): Null Models. Training on random labels (grey) maintains high BDM, ruling out optimization artifacts. Crucially, shuffling weights (purple star) at the grokking point restores maximal complexity, proving BDM measures topological structure, not just magnitude.
Figure 7: Granular analysis and sensitivity checks. Left (B.3): Layer-wise BDM Dynamics. Simplification is non-uniform; intermediate layers (red/white) show the most dramatic complexity reduction, consistent with the formation of the core algorithmic circuit. Right (B.4): Hyperparameter Sensitivity. Varying quantization levels ( ) and block sizes ( to ) shifts the absolute scale but preserves the inflection point and trajectory shape. The BDM signal is robust to these ad hoc choices.
Figure 8: Universality across model depths. Left (B.5a): Dynamics Comparison. We track Test Accuracy (blue) and BDM Complexity (red) for both the deep 48-layer model (solid) and the shallow 2-layer baseline (dashed). The phase transition signature—a synchronized accuracy jump and complexity collapse—is identical in both regimes, confirming that grokking is not an artifact of depth. Right (B.5b): Functional Topology. Layer-wise causal mediation analysis of the deep model reveals a ”U-shaped” profile, where the vast majority of intermediate layers (approx. layers 4-44) are functionally silenced. This confirms that the deep model actively discovers and implements the minimal circuit (functionally equivalent to the shallow baseline) via complexity regularization, validating the ”bypass” not as an artifact, but as a mechanism of parsimony.
Figure 9: Stability and specificity of causal interventions. Left (B.6a): Cross-Seed Stability. We track the Causal Mediation Score (CMS) of the top-5 identified functional heads across 5 independent training runs. The low variance (tight boxplots) confirms that these specific heads (or their architectural equivalents) consistently emerge as the primary drivers of the computation. Right (B.6b): Positional Specificity. We applied the patching intervention at every token position in the sequence ‘[x, op, y, =, ans]‘. The causal effect is strictly localized to the final ”Answer” token, confirming that our identified circuit is responsible for the final readout/calculation step, rather than merely disrupting upstream information flow.
Figure 10: Distribution of causal effects across contexts. We measured the logit difference (improvement in correct answer probability) induced by restoring the identified circuit on randomly sampled arithmetic problems. The distribution for the circuit (green) is shifted significantly positive, showing a consistent mechanistic effect, distinct from the zero-centered noise of random controls (grey).
Figure 11: Consensus with standard matrix measures. We track the normalized evolution of BDM (red) alongside Nuclear Norm (blue) and Stable Rank (green). All metrics show a synchronized collapse at the grokking point. Since Nuclear Norm and Stable Rank are strictly mathematical properties of the weight spectrum with no tunable parameters, their agreement with BDM confirms that the observed complexity reduction is an objective physical phenomenon, not an artifact of our measurement tool.
Figure 12: Structural simplification predicts functional importance. Each point represents a Transformer layer. We observe a strong positive correlation ( ) between a layer’s complexity reduction ( BDM) and its causal contribution to the task (CMS). Layers that ”grok” (simplify) the most are precisely those that drive the model’s performance. This suggests that low-complexity structure is not merely a byproduct, but a necessary condition for functional generalization in this domain.
Figure 13: Scaling laws and theoretical calibration. Left (B.10a): Modulus Sweep. We track the grokking time as a function of modulus size . The data follows a clear power law , consistent with the hypothesis that the model must search a hypothesis space that grows polynomially with before finding the sparse generalized solution. Right (B.10b): SFM Calibration. We plot the theoretically predicted critical sample size from the SFM equations against the empirically observed training steps . The strong linear correlation confirms that the SFM serves as a valid quantitative proxy for Transformer training dynamics, effectively bridging the gap between statistical physics theory and deep learning practice.
Figure 14: Robustness to regularization and baselines. Left (C.1): Weight Decay (WD) Sweep. A distinct phase transition in the Gini coefficient is observed across weight decay settings, including (light blue). Stronger WD ( , dark blue) typically increases the magnitude of spectral concentration, so our claim is robustness of the transition pattern rather than strict WD-invariance of absolute values. Right (C.2): Comparison with Baselines. The learned solution (red) exhibits higher spectral structure than both random Gaussian initialization (dotted black) and fixed random Fourier features (dashed grey). Shuffling the learned weights (purple star) destroys this structure, indicating that the metric captures topological alignment rather than only value distribution.
Figure 15: Universality across Moduli. The spectral Gini coefficient exhibits the same phase transition behavior across different prime moduli . The delay in onset time for larger is consistent with the increased search complexity predicted by the SFM theory.
Figure 16: Beyond visualization: Quantitative verification of the manifold loop. Left (C.4): Topological Persistence. We track the lifetime of the most persistent 1-dimensional homological feature (Betti-1 loop). The sharp rise in persistence (red) coinciding with grokking confirms that a true topological loop forms in the high-dimensional activation space, significantly distinct from random point cloud noise (grey). Right (C.5): Algebraic Alignment. The emergent geometry is not just a loop, but an algebraic loop. The circular correlation with true modular labels (blue) approaches 1.0, while the homomorphism error (orange)—measuring the deviation from linearity in the phase space—collapses to zero. This confirms the representation is isomorphic to the cyclic group .
More in this topic
Grokking
Grokking is delayed generalization in which neural networks memorize training data before abruptly learning a rule that works on unseen examples.
See all 77 papers