You train a Transformer on modular arithmetic—adding two numbers and taking the remainder.
Within five minutes, the model memorizes every example in the training set. 100% accuracy on everything it has seen. But on new problems—unseen number combinations under the same modular operation—it scores only 30%, essentially random guessing.
You keep training. Five thousand steps. Eight thousand. Ten thousand.
Nothing changes.
Then around twelve thousand steps, it's as if someone flipped a switch: accuracy starts climbing from 30%. Fast. Within forty steps it hits 95%. The model has "understood" how modular arithmetic works. Not memorization anymore—actual understanding.
This phenomenon has a cool name: grokking—a Martian word from Robert Heinlein's *Stranger in a Strange Land*, where "grok" means "to understand completely and intimately." ML researchers borrowed it to describe this eerie transition from rote memorization to sudden comprehension.
What's eerie isn't that the model eventually generalizes—that's what we want. It's that it hovers between understanding and memorization for ten thousand steps with no visible change, then transitions in an instant.
🕰️ What happens during those ten thousand steps?
Why doesn't it work like: the model memorizes first, then—once it has fully memorized the training data—starts learning to understand *as well*? 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 what conditions trigger "sudden understanding." Worse, you may not be able to predict when a model will grok—or 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 eating away connections that help memorization but not generalization. Some say it's waiting for certain features to emerge. Some say it's waiting for a sparse subnetwork—a "winning lottery ticket"—to be found.
Hidajat, Stoll, and An recently posted a paper on arXiv (2605.15787) offering a different angle. They argue that 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 or hidden dimensions you use—cannot recover that information. Drop it once, and it's gone forever.
This sounds trivial, but it has deep consequences.
Hidajat et al. model attention as Bayesian posterior estimation—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 receive enough mass—no critical token can be lost.
This yields two conditions:
1. MLP capacity must be neither too small nor too large. A too-small MLP can't memorize enough patterns; a too-large one memorizes with excess parameters, actively hindering generalization. This is consistent with earlier theory.
2. Attention must place sufficient mass on every informative token. This condition was absent from earlier theories—because it only exists in attention mechanisms.
💥 The problem is explaining away
Their core argument: in early training, the MLP is capable of direct memorization—by learning task-irrelevant "shortcut features" that fit the training data. Once the MLP drives cross-entropy loss to near zero, the whole model "feels" done. Gradient signals become tiny. Attention stops receiving useful feedback.
Worse, since the MLP has already "explained" all the outputs, attention loses its incentive to learn the task structure—a phenomenon from causal inference called "explaining away."
So when can attention resume learning? Only after weight decay erodes the MLP's memorization traces. Once memory is weakened, 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 an elegant experiment: they added a KL divergence term to the objective, directly pulling attention's distribution toward a prior—bypassing the vanished gradient signal.
Result: the grokking delay shrank dramatically, following a clean scaling law—double the intervention strength, halve the waiting time.
This is practical. If you know a model is "stuck in the memorization phase" because attention isn't receiving gradients, you don't have to passively wait for weight decay to erode memory. You can just push it.
🤷 What I don't know
A few honest confessions:
First, how does this "Bayesian lottery ticket" idea relate to the famous Lottery Ticket Hypothesis? The paper's title deliberately echoes it. The Lottery Ticket Hypothesis holds that networks contain small, independently trainable subnetworks, and grokking is the process of finding them. This paper says attention additionally needs a structural condition. Are these explanations complementary or competing? Even after reading the full paper, I couldn't fully tell whether this is a layered account or a rival hypothesis.
Second, I noticed first author Kai Hidajat also posted another paper in the same arXiv batch—"Martingale Neural Operators" (2605.15806)—on learning stochastic marginals via Doob-Meyer decomposition. That's a broad research range. Grokking and martingale theory from the same researcher—I'm curious whether shared mathematical tools connect the two lines of work, but I lack the background to judge.
Third, their experiments are on algorithmic sequence tasks (modular arithmetic, parity). These tasks have clean structure. Does grokking on fuzzier tasks—like natural language syntax—follow the same delayed mechanism? No idea. Natural language dependency graphs are far messier than modular arithmetic, and attention's "Bayesian inference" may behave completely differently in such noisy environments.
❓ One thing is clear
The paper decomposes grokking—one of the most puzzling phenomena in ML—into two separable conditions: a "Goldilocks" MLP capacity condition and a Bayesian structural condition on attention. And it explains why the two are separated in time—not because one requires long training, but because the first "explains away" the second.
If you've ever seen a Transformer's curve sitting flat for thousands of steps and then suddenly understanding in one—that's what this is about. This paper tells you why it waited. And 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.