What is Grokking?
Grokking is a striking phenomenon in machine learning: a model first memorizes the training data, then—after a long delay—suddenly generalizes. When a Transformer is trained on a small dataset like modular arithmetic, it quickly memorizes every training example (100% training accuracy) while still performing near chance (~30%) on unseen combinations. Training appears stagnant for thousands of steps, until around step 12,000, when generalization accuracy jumps abruptly from ~30% to over 95% within a few dozen steps. The name comes from Robert Heinlein's Martian word "grok," meaning to understand fully and intuitively.
The puzzle is not that the model eventually generalizes, but why it waits so long—hovering between memorization and understanding for ~10,000 steps with no visible progress, then transitioning almost instantaneously.
Previous Explanations
Researchers had proposed several accounts of the grokking delay, none fully accounting for attention's peculiarities:
- Norm minimization: large-magnitude weights support memorization; weight decay gradually shrinks them, pushing the model toward smaller-norm, generalizing solutions (consistent with the Goldilocks capacity constraint).
- Feature emergence: the model initially relies on simple, memorization-friendly features; more abstract, generalizable features emerge over training.
- Lottery Ticket Hypothesis: random initialization hides sparse trainable subnetworks ("winning tickets"); grokking occurs when training finally discovers a generalizing parameter subset.
- The grokking wait time drops substantially.
- A clean scaling law emerges: doubling the intervention strength roughly halves the waiting time.
- Relation to the Lottery Ticket Hypothesis: the paper's title ("Bayesian Lottery Tickets") invites comparison. Lottery tickets ask *which parameters* generalize; this work asks *which information* is used. Whether they are complementary or competing remains open.
- Task scope: experiments cover algorithmic tasks (modular addition, parity) with clean discrete dependency structures. Noisier tasks like natural language syntax may behave differently.
- Other architectures: grokking also occurs in RNNs and MLPs, which lack explicit attention; whether the same mechanism applies is unknown.
All of these overlook attention's role as an information bottleneck: if attention drops a critical piece of information at some layer, no later computation can recover it.
Attention as a Gambler
Hidajat et al.'s new paper reframes attention as a gambler placing bets on which input tokens matter. In Bayesian terms, attention implicitly infers the task's dependency structure, assigning each token a posterior-like weight. If attention drops an informative token, no amount of subsequent finite computation can compensate—the information is lost forever, unlike in convolutional or fully connected networks where later layers can recombine signals.
Thus, generalization requires attention to place sufficient mass on every informative token.
Two Separable Conditions for Generalization
1. Goldilocks constraint on MLP capacity: capacity must be moderate—too little cannot fit the data; too much enables memorization via shortcut features. 2. Bayesian structural condition on attention: attention must correctly infer the task's dependency graph, giving nonzero posterior mass to *all* informative nodes.
These conditions are separable, and the paper's thesis is that delayed generalization is delayed structure inference.
Explaining Away: Why Understanding Waits for Memory to Fade
Early in training, the MLP has enough capacity to memorize the data, driving cross-entropy loss near zero. Gradients become tiny, so attention receives almost no useful feedback about task structure. This is the causal-inference phenomenon of explaining away: because the MLP's memorization already "explains" the outputs, attention's structural inference is suppressed.
Only when weight decay erodes the memorized solution does loss rise and gradient signal reappear—allowing attention to finally learn the missed dependencies. Grokking waits because it waits for memorization to fade. This explains why grokking time is inversely related to weight-decay strength: stronger decay erodes memory faster, so attention gets feedback sooner.
Breaking the Wait: A Scaling Law for Structural Intervention
The authors add a KL-divergence term to the objective, pulling attention toward a structural prior (e.g., encouraging more uniform attention to avoid dropping tokens). Results:
This KL intervention bypasses the explaining-away trap by giving attention direct structural guidance regardless of gradient weakness.
Open Questions and Limitations
Summary
The paper decomposes grokking into two separable conditions—a Goldilocks capacity constraint and a Bayesian structural condition on attention—identifying delayed generalization as delayed structure inference, suppressed early by memorization via explaining away. A KL structural intervention shortens the wait with a predictable scaling law, showing that grokking's delay is not inevitable once its mechanism is understood.