You train a Transformer to do modular arithmetic—adding two numbers, then taking the remainder.
Within five minutes, the model memorizes every example in the training set. It recites the correct answer to every question, 100% accuracy. But on new problems—unseen number combinations under the same modular arithmetic—it only gets about 30% right, essentially random guessing.
You keep training. Five thousand steps. Eight thousand. Ten thousand.
Nothing changes.
Then, around step twelve thousand, it's as if someone flipped a switch: accuracy starts climbing from 30%. Fast. Within forty steps it hits 95%. It "understands" modular arithmetic. Not memorization anymore—actual understanding.
This phenomenon has a cool name: grokking—from the Martian word in Robert Heinlein's *Stranger in a Strange Land*, where "grok" means "to understand completely, to merge with." ML researchers borrowed the term to describe this eerie transition from rote memorization to sudden understanding.
What's eerie isn't that it eventually understands—we want that. It's that it sits between memorization and understanding for ten thousand steps, nothing happens, and then the transition happens in an instant.
🕰️ What happens during those ten thousand steps?
Why isn't it: the model memorizes first, and then—once it has memorized all training data—it starts learning to "also" understand? Why does it wait so long?
This isn't a boring question. If you don't know why the model waits, you don't know the conditions for sudden understanding. Worse, you may not be able to predict when—or whether—it will grok on new problems.
Many have tried to explain this waiting period. Some say weight decay is at work—slowly eating away connections that only help memorization, not understanding. Some say it's waiting for certain features to emerge. Some say it's waiting for some sparse subnetwork—a "winning lottery ticket"—to be found.
Hidajat, Stoll, and An recently posted a paper on arXiv (2605.15787) proposing a different angle: the attention mechanism itself is what's holding things back.
👀 Attention is a gambler
The key insight: if attention drops an informative token at some layer, all computation after that layer—no matter what MLPs or hidden dimensions you use—cannot recover that information. Dropped once, lost forever.
This sounds trivial, but it has deep consequences.
Hidajat et al. model attention as a Bayesian posterior estimate—attention is "guessing" what the task's dependency graph looks like. Every token is either "informative" or "noise," and attention must assign a weight to each. Correct generalization requires: every informative token must receive enough probability mass—no critical token can be dropped.
This yields two conditions:
1. MLP capacity must be neither too large nor too small—a too-small MLP can't memorize enough patterns; a too-large one devotes too many parameters to rote memorization, blocking generalization. This is consistent with earlier theory.
2. Attention must place sufficient mass on every informative token. This condition wasn't considered in prior theories—because it only exists in attention.
💥 The problem is explaining away
Their core argument: early in training, the MLP is capable of outright memorization—fitting the training data by learning task-irrelevant "shortcut features." Once the MLP pushes cross-entropy loss near zero, the whole model "feels" done. Gradient signals become tiny. Attention stops receiving useful feedback.
Worse: because the MLP has "explained" all the outputs, attention loses its driving force to learn the task's structure—a phenomenon in causal reasoning called explaining away.
So when can attention resume learning? Only after weight decay erodes part of the MLP's memorized traces. As memory weakens, loss rises again, gradient signals reappear, and attention finally gets feedback.
This explains why grokking takes so long: it's not waiting for attention to improve—it's waiting for memorization to decay.
🔧 Can we skip the wait?
Hidajat et al. ran a clever experiment: they added a KL divergence term to the objective—directly pulling the attention distribution toward a prior, bypassing the vanishing-gradient problem.
Result: the grokking delay shrank dramatically, with a clean scaling law—every doubling of intervention strength halved the waiting time.
This is very practical. If you know the model is "stuck in the memorization phase" because attention gets no gradient, you don't need to passively wait for weight decay. You can just push it.
🤷 What I don't know
A few honest caveats:
1. What's the relationship between this "Bayesian lottery ticket" idea and the famous Lottery Ticket Hypothesis? The paper's title—"Bayesian Lottery Tickets"—seems deliberately in dialogue with it. The Lottery Ticket Hypothesis says networks hide small, independently trainable subnetworks, and grokking is the process of finding them. This paper says attention needs an additional structural condition. Are these explanations complementary or competing? Even after reading the full paper, I couldn't tell whether this is a layered account or an opposing hypothesis.
2. First author Kai Hidajat posted another paper in the same arXiv batch—"Martingale Neural Operators" (2605.15806)—on learning stochastic boundaries via Doob-Meyer decomposition. That's a wide research span. The same researcher working on grokking and martingale theory simultaneously—I'm curious whether a shared mathematical tool connects the two lines, but I lack the background to judge.
3. Their experiments are on algorithmic sequence tasks (modular addition, parity). These have clean structure. Does grokking on fuzzier tasks—like natural language syntax—follow the same delay mechanism? Completely unknown. Natural language dependency graphs are far messier than modular arithmetic, and attention's "Bayesian inference" may behave totally differently in such noisy environments.
❓ But one thing is clear
This paper takes one of the most perplexing ML phenomena—grokking—and decomposes it into two separable conditions: a "Goldilocks" condition on MLP capacity, and a Bayesian structural condition on attention. Then it explains why the two conditions are separated in time—not because one needs long training, but because the first "explains away" the second.
If you've ever seen a Transformer do nothing for four thousand steps and then suddenly understand in a single step—you know what I'm talking about. This paper tells you why it waits. And tells you what to do if you don't want to wait.
---
References
1. Hidajat, K., Stoll, S., & An, J. (2026). *Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets*. arXiv:2605.15787 [cs.LG]. https://arxiv.org/abs/2605.15787
2. Power, A., Burda, Y., Edwards, H., Babuschkin, I., & Misra, V. (2022). *Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets*. arXiv:2201.02177.
3. Frankle, J., & Carbin, M. (2019). *The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks*. ICLR 2019.
4. Nanda, N., Chan, L., Lieberum, T., Smith, J., & Steinhardt, J. (2023). *Progress Measures for Grokking via Mechanistic Interpretability*. ICLR 2023.
5. Vardi, G. (2023). *On the Interplay Between Memorization and Generalization in Neural Networks*. Journal of Machine Learning Research, 24(1), 1-48.