mlsanityData to 7 Oct 2026

Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

arXiv 2605.20441 ↗First posted 19 May 2026cs.LG cs.AI cs.NEcited by 2

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

  1. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  2. 2022Emergent Abilities of Large Language Modelscited by 2,483
  3. 2018The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networkscited by 2,250
  4. 2022In-context Learning and Induction Headscited by 930
  5. 2022Toy Models of Superpositioncited by 928
  6. 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 844
  7. 2023Progress measures for grokking via mechanistic interpretabilitycited by 822
  8. 2019Are Sixteen Heads Really Better than One?cited by 803
  9. 2013Exact solutions to the nonlinear dynamics of learning in deep linear neural networkscited by 748
  10. 2020Prevalence of Neural Collapse during the terminal phase of deep learning trainingcited by 713
  11. 2022Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetscited by 599
  12. 2023Are Emergent Abilities of Large Language Models a Mirage?cited by 515
38 more
  1. 2016Exponential expressivity in deep neural networks through transient chaoscited by 260
  2. 2018A mathematical theory of semantic development in deep neural networkscited by 182
  3. 2022Omnigrok: Grokking Beyond Algorithmic Datacited by 142
  4. 2023The Quantization Model of Neural Scalingcited by 137
  5. 2020Predicting trends in the quality of state-of-the-art neural networks without access to training or testing datacited by 113
  6. 2023Explaining grokking through circuit efficiencycited by 105
  7. 2023Grokking as the Transition from Lazy to Rich Training Dynamicscited by 102
  8. 2023Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokkingcited by 76
  9. 2023Why Do We Need Weight Decay in Modern Deep Learning?cited by 70
  10. 2024Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformerscited by 46
  11. 2024How to set AdamW's weight decay as you scale model and dataset sizecited by 42
  12. 2024Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficientcited by 39
  13. 2024Artificial Kuramoto Oscillatory Neuronscited by 37
  14. 2022Grokking phase transitions in learning local rules with gradient descentcited by 28
  15. 2024An exactly solvable model for emergence and scaling laws in the multitask sparse parity problemcited by 21
  16. 2026There Will Be a Scientific Theory of Deep Learningcited by 21
  17. 2022SGD and Weight Decay Secretly Minimize the Rank of Your Neural Networkcited by 17
  18. 2023From Coupled Oscillators to Graph Neural Networks: Reducing Over-smoothing via a Kuramoto Model-based Approachcited by 15
  19. 2025Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokkingcited by 14
  20. 2026Low-Dimensional and Transversely Curved Optimization Dynamics in Grokkingcited by 10
  21. 2025Is Grokking a Computational Glass Relaxation?cited by 8
  22. 2025The Geometry of Grokking: Norm Minimization on the Zero-Loss Manifoldcited by 8
  23. 2026The Geometry of Multi-Task Grokking: Transverse Instability, Superposition, and Weight Decay Phase Structurecited by 8
  24. 2025What Can Grokking Teach Us About Learning Under Nonstationarity?cited by 7
  25. 2026To Grok Grokking: Provable Grokking in Ridge Regressioncited by 6
  26. 2026Early-Warning Signals of Grokking via Loss-Landscape Geometrycited by 6
  27. 2026The Norm-Separation Delay Law of Grokking: A First-Principles Theory of Delayed Generalizationcited by 6
  28. 2026A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularizationcited by 5
  29. 2026The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topologycited by 4
  30. 2026Spectral Edge Dynamics: An Analytical-Empirical Study of Phase Transitions in Neural Network Trainingcited by 4
  31. 2026Grokking as Dimensional Phase Transition in Neural Networkscited by 4
  32. 2026Model Capacity Determines Grokking through Competing Memorisation and Generalisation Speedscited by 4
  33. 2026Grokking as a Falsifiable Finite-Size Transitioncited by 3
  34. 2026From Density Matrices to Phase Transitions in Deep Learning: Spectral Early Warnings and Interpretabilitycited by 3
  35. 2026The Long Delay to Arithmetic Generalization: When Learned Representations Outrun Behaviorcited by 3
  36. 2026Dimensional Criticality at Grokking Across MLPs and Transformerscited by 3
  37. 2023A Waddington landscape for prototype learning in generalized Hopfield networkscited by 2
  38. 2026Late-Stage Generalization Collapse in Grokking: Detecting anti-grokking with Weightwatchercited by 2

50 most cited of 63 All 63 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2025From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics
  2. 2026Stability and Generalization in Looped Transformers
  3. 2026Invariant Learning Dynamics of Transformers in Inductive Reasoning Tasks
  4. 2025Two-Scale Latent Dynamics for Recurrent-Depth Transformers
  5. 2025Too Big to Think: Capacity, Memorization, and Generalization in Pre-Trained Transformers
  6. 2025From Memorization to Reasoning in the Spectrum of Loss Curvature
  7. 2023AttMEMO : Accelerating Transformers with Memoization on Big Memory Systems
  8. 2026Mixture of Chapters: Scaling Learnt Memory in Transformers