You train a Transformer to do modular addition — adding two numbers, then taking the remainder.
Within five minutes, the model has memorized every example in the training data. It can recite the correct answer for every question it has seen: 100% accuracy. But on new problems — unseen number combinations under the same modular arithmetic — it only gets about 30% right. 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 has "understood" how modular addition works. No longer memorization — genuine understanding.
This phenomenon has a cool name: grokking — a Martian word from Robert Heinlein's novel *Stranger in a Strange Land*, where "grok" means "to understand fully and intimately." ML researchers borrowed the term to describe this eerie transition from rote memorization to sudden understanding.
What's eerie about it? Not that it eventually understands — that's obviously what we want. What's eerie is that it lingers between memorization and understanding for ten thousand steps with nothing happening, then transitions in an instant.
🕰️ What happens during those ten thousand steps?
Why isn't it: the model memorizes first, then — once it has fully memorized all training data — starts to "also" learn to understand? Why does it wait so long?
This is not a boring question. If you don't know why the model waits, you don't know the conditions under which it "suddenly understands." Worse, you may be unable to predict when it will grok — or even whether it will grok at all on new problems.
Many have tried to explain this waiting period. Some say weight decay is at work — slowly gnawing away at connections that help memorization but not generalization. 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 the bottleneck.
👀 Attention is a gambler
The key insight: if attention drops an informative token at some layer, then all computation after that layer — no matter what MLPs you use or how large the hidden dimension — cannot recover that information. Drop it once, and it's gone forever.
That sounds like a truism, 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. Each token is either "informative" or "noise," and attention must assign a weight to each. Correct generalization requires that every informative token receives enough attention mass — no critical token can be dropped.
This yields two conditions.
The first is that MLP capacity must be neither too small nor too large — a too-small MLP can't memorize enough patterns; a too-large one devotes too many parameters to rote memorization, which hinders generalization. This is consistent with some earlier theory.
The second condition is that attention must place sufficient mass on every informative token. This condition was not considered in prior theory — because it exists only in attention mechanisms.
💥 The problem is explaining away
Their core argument: early in training, the MLP is capable of direct memorization — fitting the training data via task-irrelevant "shortcut features." Once the MLP drives cross-entropy loss near zero, the whole model "feels" it has done well. Gradient signals become tiny. Attention stops receiving useful feedback.
Worse, since the MLP has already "explained" all the outputs, attention loses its drive to learn task structure — a phenomenon known in causal reasoning as "explaining away."
So when can attention learn again? Only after weight decay has partially eroded the MLP's memorization traces. Weakened memorization means loss rises again, gradient signals reappear, and attention can finally receive feedback.
This explains why grokking takes so long: it isn't waiting for attention to improve — it's waiting for the memorization to be degraded.
🔧 Can we avoid the wait?
Hidajat et al. ran an elegant 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 — doubling the intervention strength halves the waiting time.
This is practical. If you know the model is "stuck in the memorization phase" because attention isn't receiving gradients, you don't need to passively wait for weight decay to slowly erode the memorization. You can give it a push directly.
🤷 What I don't know
A few things I must admit.
First, what is the relationship between this "Bayesian lottery" idea and the famous Lottery Ticket Hypothesis? The paper's title uses "Bayesian Lottery Tickets," seemingly in deliberate dialogue with it. The Lottery Ticket Hypothesis holds that small, independently trainable subnetworks are hidden inside neural networks, and grokking is the process of finding them. This paper says attention also requires an additional structural condition. Are these explanations complementary or competing? After reading the full text, I still couldn't fully determine whether this is a layered explanation or an opposing hypothesis.
Second, I noticed the first author, Kai Hidajat, also published another paper in the same arXiv batch — "Martingale Neural Operators" (2605.15806) — on learning stochastic martingales 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 thread connects the two lines, but I lack the background to judge.
Third, their experiments were done on algorithmic sequence tasks (modular addition, parity). These tasks have clean structure. Does grokking on fuzzier tasks — like syntactic structure in natural language — follow the same delay mechanism? Completely unknown. Natural language dependency graphs are nowhere near as clean as modular addition, and attention's "Bayesian inference" may behave entirely differently in such noisy environments.
❓ But one thing is clear
This paper breaks grokking — one of the most perplexing phenomena in ML today — into two separable conditions: a "Goldilocks" condition on MLP capacity, and a Bayesian structural condition on attention. And it explains why these two conditions are separated in time — not because one requires long training, but because the first "explains away" the second.
If you've ever seen the curve of a Transformer doing nothing for four thousand steps and then suddenly understanding in one — you know what I'm talking about. This paper tells you why it waited. And tells you what you can do if you'd rather not 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.