Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics
Not on the map. Why?
Abstract (arXiv)
Transformers trained on modular arithmetic exhibit sharp transitions between memorization, generalization, and collapse. We show that weight decay acts as a scalar empirical control parameter for these regimes, and introduce two cheap online diagnostics, mean pairwise attention-head cosine similarity and entropy standard deviation, that track training dynamics from attention activations alone and complement loss-landscape diagnostics at lower compute cost. Across eleven experimental conditions and three model scales (0.82M to 85M parameters), the weight-decay axis separates memorization, developmental grokking, and collapse. A near-transition logistic fit localizes the memorization-to-developmental boundary at λc=0.0158 (95% CI [0.0109, 0.0200], N=210); a power-law fit gives an empirical exponent nu=0.757 (CI [0.725, 0.799]). Reference exponents nu=1/2 and 3D Ising nu ≈ 0.63 lie outside this empirical CI under our four-bin grid, so we report nu as empirical and defer universality-class identification to denser finite-size-scaling work. A horizon-matched multi-task replication (n=280, four modular operations) preserves the weight-decay control pattern; a paired attention-head re-initialization experiment at λ=0.05 changes Phase-2 amplitude (Cohen's d=-1.190, n=10, pt=4.5 × 10-3), while matched weight-norm clipping does not. Three cross-architecture probes (4L MLP, 4L LSTM, and 4L Mamba; each n=70) replicate the weight-decay-controlled transition with architecture-specific λc values. Main diagnostic claims are scoped to modular arithmetic in small transformer attention models; the non-attention experiments are scope probes, and architecture-wide, language-model, and universality-class claims are out of scope.
Key facts. Cited by 2 papers in the map (data through 7 Oct 2026). Builds on 63 of its 68 references that are in the map. First posted on arXiv 19 May 2026.
Builds on
63 of 68 references.
References · 63
- 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
- 2022Emergent Abilities of Large Language Modelscited by 2,483
- 2018The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networkscited by 2,250
- 2022In-context Learning and Induction Headscited by 930
- 2022Toy Models of Superpositioncited by 928
- 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 844
- 2023Progress measures for grokking via mechanistic interpretabilitycited by 822
- 2019Are Sixteen Heads Really Better than One?cited by 803
- 2013Exact solutions to the nonlinear dynamics of learning in deep linear neural networkscited by 748
- 2020Prevalence of Neural Collapse during the terminal phase of deep learning trainingcited by 713
- 2022Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetscited by 599
- 2023Are Emergent Abilities of Large Language Models a Mirage?cited by 515
38 more
- 2016Exponential expressivity in deep neural networks through transient chaoscited by 260
- 2018A mathematical theory of semantic development in deep neural networkscited by 182
- 2022Omnigrok: Grokking Beyond Algorithmic Datacited by 142
- 2023The Quantization Model of Neural Scalingcited by 137
- 2020Predicting trends in the quality of state-of-the-art neural networks without access to training or testing datacited by 113
- 2023Explaining grokking through circuit efficiencycited by 105
- 2023Grokking as the Transition from Lazy to Rich Training Dynamicscited by 102
- 2023Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokkingcited by 76
- 2023Why Do We Need Weight Decay in Modern Deep Learning?cited by 70
- 2024Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformerscited by 46
- 2024How to set AdamW's weight decay as you scale model and dataset sizecited by 42
- 2024Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficientcited by 39
- 2024Artificial Kuramoto Oscillatory Neuronscited by 37
- 2022Grokking phase transitions in learning local rules with gradient descentcited by 28
- 2024An exactly solvable model for emergence and scaling laws in the multitask sparse parity problemcited by 21
- 2026There Will Be a Scientific Theory of Deep Learningcited by 21
- 2022SGD and Weight Decay Secretly Minimize the Rank of Your Neural Networkcited by 17
- 2023From Coupled Oscillators to Graph Neural Networks: Reducing Over-smoothing via a Kuramoto Model-based Approachcited by 15
- 2025Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokkingcited by 14
- 2026Low-Dimensional and Transversely Curved Optimization Dynamics in Grokkingcited by 10
- 2025Is Grokking a Computational Glass Relaxation?cited by 8
- 2025The Geometry of Grokking: Norm Minimization on the Zero-Loss Manifoldcited by 8
- 2026The Geometry of Multi-Task Grokking: Transverse Instability, Superposition, and Weight Decay Phase Structurecited by 8
- 2025What Can Grokking Teach Us About Learning Under Nonstationarity?cited by 7
- 2026To Grok Grokking: Provable Grokking in Ridge Regressioncited by 6
- 2026Early-Warning Signals of Grokking via Loss-Landscape Geometrycited by 6
- 2026The Norm-Separation Delay Law of Grokking: A First-Principles Theory of Delayed Generalizationcited by 6
- 2026A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularizationcited by 5
- 2026The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topologycited by 4
- 2026Spectral Edge Dynamics: An Analytical-Empirical Study of Phase Transitions in Neural Network Trainingcited by 4
- 2026Grokking as Dimensional Phase Transition in Neural Networkscited by 4
- 2026Model Capacity Determines Grokking through Competing Memorisation and Generalisation Speedscited by 4
- 2026Grokking as a Falsifiable Finite-Size Transitioncited by 3
- 2026From Density Matrices to Phase Transitions in Deep Learning: Spectral Early Warnings and Interpretabilitycited by 3
- 2026The Long Delay to Arithmetic Generalization: When Learned Representations Outrun Behaviorcited by 3
- 2026Dimensional Criticality at Grokking Across MLPs and Transformerscited by 3
- 2023A Waddington landscape for prototype learning in generalized Hopfield networkscited by 2
- 2026Late-Stage Generalization Collapse in Grokking: Detecting anti-grokking with Weightwatchercited by 2
50 most cited of 63 All 63 (JSON)
Built on by
Citing papers per year: 2026: 2.
2 papers.
Similar, not cited
Close in wording, no citation either way
- 2025From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics
- 2026Stability and Generalization in Looped Transformers
- 2026Invariant Learning Dynamics of Transformers in Inductive Reasoning Tasks
- 2025Two-Scale Latent Dynamics for Recurrent-Depth Transformers
- 2025Too Big to Think: Capacity, Memorization, and Generalization in Pre-Trained Transformers
- 2025From Memorization to Reasoning in the Spectrum of Loss Curvature
- 2023AttMEMO : Accelerating Transformers with Memoization on Big Memory Systems
- 2026Mixture of Chapters: Scaling Learnt Memory in Transformers