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

Mar 26, 2026

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

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.

Shalima Binta Manir, Anamika Paul Rupa

Howard University

Grokking Neural Network Training Dynamics Neural Network Generalization Neural Network Optimization

Ask PDF

Feb 18, 2026

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Grokking -- the delayed transition from memorization to generalization in small algorithmic tasks -- remains poorly understood. We present a geometric analysis of optimization dynamics in transformers trained on modular arithmetic. PCA of attention weight trajectories reveals that training evolves predominantly within a low-dimensional execution subspace, with a single principal component capturing 68-83% of trajectory variance. To probe loss-landscape geometry, we measure commutator defects -- the non-commutativity of successive gradient steps -- and project them onto this learned subspace. We find that curvature grows sharply in directions orthogonal to the execution subspace while the trajectory remains largely confined to it. Importantly, curvature growth consistently precedes generalization across learning rates and hyperparameter regimes, with the lead time obeying a power law in the grokking timescale. Causal intervention experiments show that motion along the learned subspace is necessary for grokking, while artificially increasing curvature is insufficient. Together, these results support a geometric account in which grokking reflects escape from a metastable regime characterized by low-dimensional confinement and transverse curvature accumulation. All findings replicate across this learning-rate range, a qualitatively different slow regime (lr=5e-5, wd=0.1, 3 layers), and three random seeds, though alignment dynamics differ quantitatively between regimes. Causal intervention experiments establish that orthogonal gradient flow is necessary but not sufficient for grokking: suppressing it prevents generalization with a monotonic dose-response across four operations, while artificially boosting curvature defects has no effect.

Yongzhong Xu

Grokking Neural Network Training Dynamics Gradient Descent Dynamics Loss Landscape Geometry

Ask PDF

Jul 5, 2026

Structure-Specific Representational Priors Causally Control the Grokking Delay

Grokking -- generalization long after training-set interpolation -- has been accelerated by structure-agnostic interventions (gradient filtering, weight-norm clamping, geometric penalties). Whether the delay specifically measures the time to form task-structured representations has remained observational. We test it causally by injecting representational priors of varying content into a one-layer transformer learning modular addition, via a supervised-contrastive loss whose positives encode (i) the task's true structure (

(a+b) \bmod p

), (ii) a coherent-but-wrong sibling (

(a-b) \bmod p

), or (iii) a random partition -- all with identical loss form, strength, class sizes, and geometry. Whether generalization occurs follows a clean gradation: true 22/30 runs, sibling (same periodic features, wrong combination) 14/15, random (only memorizable) 0/20 (Fisher

p=1.3\times10^{-7}

). A weight-norm-matched control replaying the norm trajectory onto plain cross-entropy generalizes 0/15, ruling out the norm as mediator. Probes show structure formation precedes and predicts generalization in all runs. Only the true structure also accelerates grokking (up to

2.75\times

), but this is dose-dependent and bimodal. We then confirm the mechanism by prediction: because the acceleration is gated by a weight-norm side-effect, clamping the norm during training yields a reliable, standalone accelerator with a median

8.6\times

speedup (up to

22\times

on the fastest seeds, under 1000 epochs), growing monotonically as the norm is held lower; the residual stalls also vanish, though significant only pooled over the two mitigations run at both strengths (

0/40

6/20

p=7.7\times10^{-4}

), not per method. The grokking delay is, causally, the time to form the right representational structure -- decided at the level of features, not labels.

Gunner Levi Howe

Transformer Grokking

Ask PDF