The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology

Alper Yıldırım

cs.LGMar 5, 2026

The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology

Authors: Alper Yıldırım

Organizations: Independent Researcher

Abstract

Mechanistic interpretability typically relies on post-hoc analysis of trained networks. We instead adopt an interventional approach: testing hypotheses a priori by modifying architectural topology to observe training dynamics. We study grokking - delayed generalization in Transformers trained on cyclic modular addition (Zp) - investigating if specific architectural degrees of freedom prolong the memorization phase. We identify two independent structural factors in standard Transformers: unbounded representational magnitude and data-dependent attention routing. First, we introduce a fully bounded spherical topology enforcing L2 normalization throughout the residual stream and an unembedding matrix with a fixed temperature scale. This removes magnitude-based degrees of freedom, reducing grokking onset time by over 20x without weight decay. Second, a Uniform Attention Ablation overrides data-dependent query-key routing with a uniform distribution, reducing the attention layer to a Continuous Bag-of-Words (CBOW) aggregator. Despite removing adaptive routing, these models achieve 100% generalization across all seeds and bypass the grokking delay entirely. To evaluate whether this acceleration is a task-specific geometric alignment rather than a generic optimization stabilizer, we use non-commutative S5 permutation composition as a negative control. Enforcing spherical constraints on S5 does not accelerate generalization. This suggests eliminating the memorization phase depends strongly on aligning architectural priors with the task's intrinsic symmetries. Together, these findings provide interventional evidence that architectural degrees of freedom substantially influence grokking, suggesting a predictive structural perspective on training dynamics.

Figures & tables

Panel of Figure 1 — Figure 1: Test accuracy across architectural configurations. (a) Standard normalizations exhibit the classic, delayed grokking profile over 100,000 epochs. (b) A zoomed-in view of the first 5,000 epochs reveals that the topologically bounded models bypass this delay entirely, exhibiting immediate and stable convergence while the baselines remain trapped at chance accuracy.

Learning Rate	Architecture	Mean Grok Epoch	Std Dev	Min	Max	Failures	Peak Acc.
$1\times 10^{-4}$	LayerNorm (Baseline)	54,160	13,490	32,800	71,600	0 / 10	100%
	RMSNorm (Baseline)	51,240	11,200	38,800	74,600	0 / 10	100%
	Bounded Sphere ( $\lambda=1.0$ )	2,400	353	1,800	3,000	0 / 10	100%
	Bounded Sphere ( $\lambda=0.0$ )	2,480	464	1,800	3,200	0 / 10	100%
	Bounded Sphere + Fourier Init ( $\lambda=1.0$ )	2,120	368	1,600	2,600	0 / 10	100%
	Bounded Sphere + Fourier Init ( $\lambda=0.0$ )	2,100	316	1,800	2,600	0 / 10	100%

Table 1: Grokking onset epoch across 10 random seeds. Spherical configurations substantially reduce convergence time at both standard and elevated learning rates.

Architecture	Mean 95% Epoch	Std Dev	Min	Max	Hits
LayerNorm (Baseline)	28,220	3,676	19,600	32,400	10 / 10
RMSNorm (Baseline)	31,360	4,056	25,600	38,600	10 / 10
Bounded Sphere ( $\lambda=1.0$ )	2,320	223	2,000	2,800	10 / 10
Bounded Sphere ( $\lambda=0.0$ )	2,200	253	1,800	2,600	10 / 10

Table 2: Generalization dynamics for modular multiplication

(a\times b)\bmod 113

across 10 random seeds. We report the epoch at which each seed first reaches 95% test accuracy as the primary generalization milestone. The bounded spherical topology accelerates generalization comparably to the addition task.

Architecture (Uniform Attention)	Mean Peak Acc.	Max Peak Acc.	Success Rate (100% Acc)
LayerNorm Baseline (WD=1.0)	100.00%	100.00%	10 / 10
Bounded Sphere ( $\lambda=1.0$ )	97.65%	100.00%	6 / 10
Bounded Sphere ( $\lambda=0.0$ )	100.00%	100.00%	10 / 10

Table 3: Peak test accuracy under Uniform Attention Ablation (Zero-Attention) across 10 random seeds. Ablating data-dependent routing allows normalized models to completely bypass the grokking delay.

Panel of Figure 2 — Figure 2: Training dynamics under Uniform Attention Ablation (Seed 2). (a) The LayerNorm baseline immediately generalizes when data-dependent routing is removed, bypassing grokking. (b) Imposing a spherical constraint while retaining weight decay introduces optimization instability, delaying generalization. (c) The Bounded Sphere without weight decay ( $\lambda=0.0$ ) stabilizes the geometry, resulting in immediate and flawless generalization.

Architecture	Mean Grok Epoch	Std Dev	Min	Max	Failures
LayerNorm (Baseline)	39,900	11,942	24,400	58,400	2 / 10
RMSNorm (Baseline)	38,250	23,600	19,000	91,600	2 / 10
Bounded Sphere ( $\lambda=1.0$ )	Failed	–	–	–	10 / 10
Bounded Sphere ( $\lambda=0.0$ )	Failed	–	–	–	10 / 10

Table 4: Grokking onset epoch for the

S_{5}

permutation composition task. Metrics for baselines are calculated exclusively over successful seeds. While the baselines generalize on the majority of seeds, the strict spherical constraints fail to produce generalization within the 100,000-epoch training window.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Architecture (Uniform Attention)	Mean Peak Acc.	Max Peak Acc.	Success Rate (100%)
LayerNorm Baseline (WD=1.0)	99.12%	99.42%	0 / 5
RMSNorm Baseline (WD=1.0)	99.65%	99.94%	0 / 5
Bounded Sphere ( $\lambda=1.0$ )	2.84%	3.42%	0 / 5
Bounded Sphere ( $\lambda=0.0$ )	2.41%	2.54%	0 / 5

Appendix

Table 5: Peak test accuracy under Uniform Attention Ablation on

S_{5}

permutation composition across 5 random seeds.

Design Choice	nGPT	Ours (Intervention A)
Goal	Faster LM pretraining	Mechanistic hypothesis test
Weight matrix normalization	All matrices, all layers	None
Hidden state update	Learnable interpolation	Residual add + hard projection
Learnable step sizes	Per-dimension $\boldsymbol{\alpha}$	None
Compensatory scaling	6 parameter families	None
Unembedding scaling	Learnable per-token $\mathbf{s}_{z}$	Fixed scalar $\tau$

Appendix

Table 6: Architectural comparison between nGPT and our Intervention A.

Explore similar work

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization
Mar 26, 2026Shalima Binta Manir, Anamika Paul RupaGrokking Neural Network Training Dynamics
Howard University
Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking
Feb 18, 2026Yongzhong XuGrokking Neural Network Training Dynamics
Structure-Specific Representational Priors Causally Control the Grokking Delay
Jul 5, 2026Gunner Levi HoweTransformer Grokking

Mar 26, 2026J/K move · Enter open · S save

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization

Shalima Binta Manir, Anamika Paul Rupa

Howard University

Grokking Neural Network Training Dynamics Neural Network Generalization Neural Network Optimization

Open paper AskPDF arXiv

Grokking the delayed transition from memorization to generalization in neural networks remains poorly understood, in part because prior empirical studies confound the roles of architecture, optimization, and regularization. We present a controlled study that systematically disentangles these factors on modular addition (mod 97), with matched and carefully tuned training regimes across models. Our central finding is that grokking dynamics are not primarily determined by architecture, but by interactions between optimization stability and regularization. Specifically, we show: (1) depth has a non-monotonic effect, with depth-4 MLPs consistently failing to grok while depth-8 residual networks recover generalization, demonstrating that depth requires architectural stabilization; (2) the apparent gap between Transformers and MLPs largely disappears (1.11

\times

delay) under matched hyperparameters, indicating that previously reported differences are largely due to optimizer and regularization confounds; (3) activation function effects are regime-dependent, with GELU up to 4.3

\times

faster than ReLU only when regularization permits memorization; and (4) weight decay is the dominant control parameter, exhibiting a narrow ``Goldilocks'' regime in which grokking occurs, while too little or too much prevents generalization. Across 3--5 seeds per configuration, these results provide a unified empirical account of grokking as an interaction-driven phenomenon. Our findings challenge architecture-centric interpretations and clarify how optimization and regularization jointly govern delayed generalization.

Abstract

Explore similar work

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Structure-Specific Representational Priors Causally Control the Grokking Delay