Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets
Not on the map. Why?
Abstract (arXiv)
Why does a Transformer that has memorized its training set wait thousands of steps before it generalizes? Existing accounts locate this delay in norm minimization, feature emergence, or the late discovery of sparse subnetworks. These explanations capture important parts of the transition, but ignore a constraint unique to attention-based models: if attention discards an informative token, no bounded downstream computation can recover it. We formalize attention as an implicit Bayesian posterior over the task dependency graph and prove that generalization requires two separable conditions: a familiar Goldilocks bound on MLP capacity, coinciding with norm-based theories of grokking, and a novel Bayesian structural condition requiring attention to place sufficient mass on every informative token. This decoupling explains delayed generalization as delayed structural inference. Early in training, the MLP memorizes through unaligned features, drives the cross-entropy loss near zero, and thereby starves attention of structural gradient. Weight decay must then erode memorization before the missing graph becomes learnable, yielding the known inverse-weight-decay delay, which we derive as a structural waiting time. We then prove that this explaining-away delay can be bypassed by a KL-based structural intervention, yielding an inverse-intervention-strength scaling law for the grokking time. Experiments on algorithmic sequence tasks isolate structure from capacity and show that this Bayesian ticket matches or outperforms lottery-ticket transfer.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 28 of its 34 references that are in the map. First posted on arXiv 15 May 2026.
Builds on
28 of 34 references.
References · 28
- 2018The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networkscited by 2,250
- 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
- 2023Progress measures for grokking via mechanistic interpretabilitycited by 822
- 2017The Implicit Bias of Gradient Descent on Separable Datacited by 649
- 2022Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetscited by 599
- 2018Gradient Descent Finds Global Minima of Deep Neural Networkscited by 572
- 2018On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transportcited by 422
- 2019Gradient Descent Maximizes the Margin of Homogeneous Neural Networkscited by 308
- 2020Gradient Starvation: A Learning Proclivity in Neural Networkscited by 262
- 2021Thinking Like Transformerscited by 182
- 2022The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networkscited by 141
- 2023Grokking as the Transition from Lazy to Rich Training Dynamicscited by 102
16 more
- 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 97
- 2023Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokkingcited by 76
- 2023Grokking of Hierarchical Structure in Vanilla Transformerscited by 53
- 2023On the Role of Attention in Prompt-tuningcited by 49
- 2025The emergence of sparse attention: impact of data distribution and benefits of repetitioncited by 25
- 2023Stack Attention: Improving the Ability of Transformers to Model Hierarchical Patternscited by 17
- 2024How Sparse Attention Approximates Exact Attention? Your Attention is Naturally nC-Sparsecited by 16
- 2024Information-Theoretic Progress Measures reveal Grokking is an Emergent Phase Transitioncited by 16
- 2025Grokking Beyond the Euclidean Norm of Model Parameterscited by 14
- 2025Provable Scaling Laws of Feature Emergence from Learning Dynamics of Grokkingcited by 14
- 2025A Theoretical Framework for Grokking: Interpolation followed by Riemannian Norm Minimisationcited by 11
- 2025From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamicscited by 9
- 2025The Geometry of Grokking: Norm Minimization on the Zero-Loss Manifoldcited by 8
- 2024Progress Measures for Grokking on Real-world Taskscited by 7
- 2025High-Dimensional Analysis of Single-Layer Attention for Sparse-Token Classificationcited by 6
- 2023Bridging Lottery Ticket and Grokking: Understanding Grokking from Inner Structure of Networkscited by 5
Built on by
Citing papers per year: 2026: 1.
1 paper.