Research map / Architectures and efficiency / RNNs and normalization

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

NeurIPS 2026 · Architectures and efficiency › RNNs and normalization · approach: Layer-wise analysis of language models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Modifies Transformer architecture before training (L2-normalized residual stream with bounded unembedding, and uniform attention) to test whether removing magnitude and routing degrees of freedom bypasses grokking delay on modular arithmetic, with S5 as a negative control.

Problem

Grokking delays generalization after memorization; the paper asks whether architectural degrees of freedom (residual magnitude and data-dependent attention routing) prolong the memorization phase, and whether architectural constraints can bypass it.

Approach

Apply a hard L2 spherical projection to the residual stream plus a normalized, fixed-temperature unembedding (Intervention A), and separately replace learned attention with uniform attention (Intervention B), applying both to Z_p arithmetic and then to non-commutative S5 as a negative control.

Key results

• Bounded Sphere + Fourier Init (lambda=0) reaches 100% test accuracy in mean 711 epochs at LR 6e-4 vs ~7,300-7,800 for LayerNorm/RMSNorm baselines.
• Under uniform attention, LayerNorm baseline and Bounded Sphere (lambda=0) both reach 100% peak test accuracy on 10/10 seeds for Z_113 addition.
• On S5, bounded spherical models fail on 10/10 seeds within 100,000 epochs while LayerNorm and RMSNorm baselines grok on 8/10 seeds (mean ~39,900 epochs).

Surprising findings

• The spherical constraint that eliminates grokking delay on Z_p completely fails on S5 (0/10 seeds), even though S5 is the simplest non-commutative task tested.
• Weight decay (lambda=1.0) introduces optimization friction under the bounded spherical topology, increasing variance and delay, and degrading Fourier circuit coherence.

Limitations

• Experiments limited to shallow one-layer Transformers on small synthetic algebraic tasks (Z_113, S5); no evaluation at scale or on natural language.
• S5 negative result tested under a single elevated learning rate (1e-3) and a limited seed count (5 for uniform attention), so conclusions about failure may depend on training setup.
• Connection to the Truong et al. norm-separation theory is interpretive and not formally validated by the experiments.

Experimental setup

Datasets: modular addition over Z_113 modular multiplication over Z_113 S5 permutation composition

Benchmarks: modular addition Z_113 modular multiplication Z_113 S5 permutation composition

Metrics: grokking onset epoch test accuracy peak accuracy fraction of variance explained (Fourier) success rate

Baselines: LayerNorm Transformer RMSNorm Transformer nGPT

Builds on: grokking (Power et al.) nGPT Softmax Collapse / StableMax Fourier circuit analysis (Nanda et al.) slingshot effect

Closest papers by meaning

• Model Capacity Determines Grokking through Competing Memorisation and Generalisation Speeds (NeurIPS 2026 · similarity 0.89)
• To Grok or not to Grok: Disentangling Generalization and Memorization on Corrupted Algorithmic Datasets (ICLR 2024 · similarity 0.89)
• Slower Generalization, Faster Memorization: A Sweet Spot in Algorithmic Learning (NeurIPS 2026 · similarity 0.89)
• Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via Grokking (NeurIPS 2025 · similarity 0.88)
• Unveiling Memorization–Generalization Coexistence: A Case Study on Arithmetic Tasks with Label Noise (NeurIPS 2026 · similarity 0.88)
• Is Grokking a Computational Glass Relaxation? (NeurIPS 2025 · similarity 0.88)