OpenAI's grokking paper never explained grokking
Three follow-up papers have since located the cause of the delay in three different places, and none of them agrees with the others.
Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets
Read the paperIn this paper we propose to study generalization of neural networks on small algorithmically generated datasets. In this setting, questions about data efficiency, memorization, generalization, and speed of learning can be studied in great detail. In some situations we show that neural networks learn through a process of "grokking" a pattern in the data, improving generalization performance from random chance level to perfect generalization, and that this improvement in generalization can happen well past the point of overfitting. We also study generalization as a function of dataset size and find that smaller datasets require increasing amounts of optimization for generalization. We argue that these datasets provide a fertile ground for studying a poorly understood aspect of deep learning: generalization of overparametrized neural networks beyond memorization of the finite training dataset.1
Perfect on the training set, wrong on everything else
Train a small transformer to fill in a table of arithmetic: given two numbers and an operation, predict the result. Show it half the table and hold out the rest. Within a few hundred update steps it reproduces every training example exactly. Test it on the withheld half, and for a long time it does no better than guessing. Then, with no change to the data, the model, or the training routine, it starts getting the withheld half right too, over a comparatively short stretch of steps that follows a very long silence.
OpenAI researchers Alethea Power, Yuri Burda, Harri Edwards, and Igor Babuschkin, working with Google's Vedant Misra, gave this jump a name in a paper posted to arXiv in January 2022 as a workshop submission. Their own definition: validation accuracy "sometimes suddenly begins to increase from chance level toward perfect generalization," long after the network has already fit the training data perfectly. They called it grokking.1 The arXiv listing itself carries no venue field, and the one workshop record found for it is dated a year before that posting, a mismatch that stays unresolved.2
That is the entire paper: a documented delay, tied to a specific task and a specific optimizer setting, with no account of why the delay ends when it does. The authors say so themselves. Their closing words float exactly one hypothesis, about flat regions of the loss landscape, and immediately flag it as untested: "It would be valuable for future work to explore this hypothesis."1
Predicting a remainder the network never saw
Take a prime p = 113. Add 87 and 61: 148. Subtract 113 once, since 148 exceeds the modulus: 35. So 87 + 61 leaves a remainder of 35 when divided by 113. A network trained on this task sees thousands of such pairs, along with the correct remainder, and has to predict the remainder on pairs held out of that training set.
Power and coauthors tested something broader: twelve binary operations, not addition alone, at a single prime modulus, p = 97. The operations ranged from division and subtraction through combinations of squares and cubes to three permutation operations on five elements. Every example ran through a two-layer, decoder-only transformer with roughly 400,000 non-embedding parameters, trained with AdamW at a learning rate of 10⁻³ and a weight decay of 1.1
What Power and coauthors actually measured
The paper's central figure trains division mod 97 with half the data held out. Training accuracy is "close to perfect" in under 1,000 optimization steps. Validation accuracy shows "very little evidence" of moving until 100,000 steps in, and takes close to 1,000,000 steps to reach the same near-perfect level training accuracy hit three orders of magnitude earlier.1
Dataset size matters more than the step count alone suggests. Near the smallest training fraction that still generalizes at all, removing just one more percent of the training data adds 40 to 50 percent to the median number of steps needed to generalize, even though the network still reaches 100 percent accuracy eventually. The number of steps needed to fit the training set itself barely moves, staying in the 1,000-to-10,000 range regardless of dataset size. Only the time to catch up on held-out data grows as the dataset shrinks.1
Weight decay changes the picture more than any other intervention tested: adding it cuts the number of training examples needed by more than half. Decaying weights toward the network's random initialization helps too, though less than decaying them toward zero.1
This is suggestive that grokking may only happen after the network's parameters are in flatter regions of the loss landscape. It would be valuable for future work to explore this hypothesis.1 Power, Burda, Edwards, Babuschkin & Misra, Appendix A.5
What the network is actually computing, and why weight decay accelerates the jump, is the subject of the published record that followed.
Addition, computed as a rotation
A year after Power's paper, five researchers, Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt, trained a transformer on the narrower addition-only task and reverse-engineered what it had learned, down to individual neurons. The paper won an oral presentation at ICLR 2023, among roughly the top quarter of that year's accepted work.3
The network's embedding layer represents each input number as coordinates on several circles at once. For a chosen frequency w, it maps a number a to the point (cos(wa), sin(wa)) on a unit circle at that frequency. Adding two numbers becomes composing two rotations around the same circle, and composing rotations is exactly what the angle-addition identity for cosine already describes:
- \cos(w_k(a+b))the frequency-k signal the correct sum should produce
- \cos(w_k a)a's position at frequency k
- \cos(w_k b)b's position at frequency k
- \sin(w_k a)a's position, a quarter-turn out of phase
- \sin(w_k b)b's position, a quarter-turn out of phase
The unembedding layer sums exactly this kind of term across five frequencies the network converges on (k = 14, 35, 41, 42, and 52, out of many it could have picked). The terms reinforce constructively at the true sum and cancel out everywhere else.3 This is not a plausible story fitted to the data after the fact. Nanda's team checked it directly: of the network's 512 hidden neurons, 433 of them, 84.6 percent, have more than 85 percent of their activity explained by a single one of those five frequencies.3 Removing every frequency except the five key ones cuts training loss 70 percent, to 7.24×10⁻⁸; keeping only the ten directions built from those five frequencies still cuts it 50 percent, to 1.19×10⁻⁷.3
Nanda's team borrows a framework called "progress measures" from a separate 2022 result on a different task entirely, sparse-parity learning, where researchers found a smooth signal of learning progress hidden beneath a training curve that looked flat right up until it wasn't.4 A progress measure is exactly that: a quantity that moves gradually and predicts a jump in visible behavior before the jump appears. Nanda's team builds four such measures for the addition task. Restricted loss zeroes out every Fourier component except the constant term and the five key frequencies. Excluded loss does the reverse: it zeroes out only the five key frequencies, measured on training data. The other two track how concentrated the per-frequency weight norms are, and the total size of the weights.3 Tracking all four across training splits it into three phases.
| Phase | Epochs | What moves |
|---|---|---|
| Memorization | 0–1 400 | Train and excluded loss fall; test and restricted loss stay high. |
| Circuit formation | 1 400–9 400 | Excluded loss rises, weight norm falls, restricted loss starts falling; test and train loss stay flat. |
| Cleanup | 9 400–14 000 | Test loss suddenly drops; weight norm falls sharply. |
Nanda's team frames the result this way: grokking is not a sudden shift but the gradual buildup of structured computation inside the weights, followed by a later cleanup that discards the memorized shortcut found first.3 That account explains what the network becomes. It does not, by itself, explain why weight decay is what drives it there, and on that question the record disagrees with itself.
Three explanations for the same delay
Weight decay speeds generalization in Power's own ablations, in Nanda's replication, and in every result described below. Why it does that is not settled. Three sets of researchers, working independently between 2022 and 2025, locate the answer in three different places, and each account genuinely conflicts with at least one of the others.
Generalization is a question of where a network's weight norm sits, not how many optimization steps it has taken.
A year of work beginning with an "effective theory" of four learning phases5 led to a specific claim. Training loss falls in an L shape against weight norm. Test loss falls in a U shape. Only a bounded "Goldilocks" band of norm values generalizes well. A network that starts outside that band first overfits, and weight decay slowly shrinks its norm back into the band. Generalization time then scales as the inverse of the decay strength. The same team reproduces grokking on image classification, sentiment analysis, and molecule-property prediction, tasks with nothing to do with modular arithmetic, though the effect is less dramatic there.6
Generalization wins a competition: it is slower to learn than memorization, but more efficient once learned.
A task that admits both a memorizing solution and a generalizing one eventually favors the generalizing one, because it produces larger output logits for the same parameter norm, and weight decay penalizes norm directly. The account predicts two behaviors neither Power's data nor Nanda's account anticipate, and the paper reports observing both. "Ungrokking" is a network that already generalizes perfectly regressing back to memorization. "Semi-grokking" is delayed generalization that stalls at partial accuracy instead of reaching it.7
Weight decay's job is arithmetic, not geometric: it keeps the network's numbers from breaking.
Without regularization, a network solving these tasks pushes some output logits large enough to exceed floating-point precision inside the softmax function, a failure the authors name Softmax Collapse: the loss numerically floors out before learning finishes. Reproducing the addition task at p = 113, they show generalization halts once collapse begins, earlier in 32-bit than 64-bit arithmetic, and that a modified activation reaches grokking with no weight decay at all.8 In their own words, the main roles of weight decay are preventing floating-point errors and preventing runaway logit scaling, not the norm-shrinking role the weight-norm account assigns it.8
None of the three papers directly rebuts the others' central claim in the sections examined here; each treats Power's original weight-decay finding as the fact to explain, but not each other's explanation of it. The record has not converged on one answer. It has produced three, standing in genuine tension, over three years.678 None of them confirms the flat-minima guess Power's own team floated in its appendix, either. The weight-norm account is about geometry, a related but distinct claim from landscape flatness; the numerical account is unrelated to either. The one hypothesis the original authors actually proposed was not tested and confirmed. It was quietly set aside.1
Where the explanation runs out
- The delay itself replicates. Power's headline curve (division mod 97, half the data withheld) and Nanda's independent addition run (p = 113, 30 percent of pairs trained on) show the same shape. Fast training accuracy, a long flat stretch on held-out data, then a jump.13
- On the narrow setup it was tested on, the rotation account is confirmed down to individual neurons and ablations, not just consistent with the training curve.3
- Weight decay reliably speeds the jump in every result examined here, regardless of which account of it turns out to be correct.18
- The rotation account is verified on one operation, one modulus, and a one-layer network, not the twelve operations, p = 97, and two-layer network Power's original claim actually covers.31
- The three accounts of weight decay contradict each other, and none has been tested head to head against the other two.678
- Power's own guess, about flat minima, was never confirmed or ruled out by the later work. It was replaced by other questions.1
Power et al. established the delay itself: documented at p up to 97, across twelve operations, with a measured dependence on weight decay and dataset size. They did not show why it happens, and said so. Nanda et al. showed why, with neuron-level ablations, for one of those twelve operations, at a different modulus, on a smaller network. That is a demonstrated cause on a narrower task, not a general one. Nothing in the record examined here shows the same rotation account holds at p = 97, on the two-layer architecture, or on the eleven other operations Power's paper reported. And the deeper question, why weight decay is what triggers the jump, is still open. Three papers published between 2022 and 2025 answer it three different ways, and disagree with each other as much as they disagree with Power's own unconfirmed guess.13678
Sources
- arXiv · Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets (Power, Burda, Edwards, Babuschkin & Misra, 2022)
- MATH-AI (ICLR 2021 Workshop) · Accepted papers listing
- arXiv · Progress Measures for Grokking via Mechanistic Interpretability (Nanda, Chan, Lieberum, Smith & Steinhardt, ICLR 2023)
- arXiv · Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit (Barak, Edelman, Goel, Kakade, Malach & Zhang, 2022)
- arXiv · Towards Understanding Grokking: An Effective Theory of Representation Learning (Liu, Kitouni, Nolte, Michaud, Tegmark & Williams, NeurIPS 2022)
- arXiv · Omnigrok: Grokking Beyond Algorithmic Data (Liu, Michaud & Tegmark, ICLR 2023)
- arXiv · Explaining Grokking Through Circuit Efficiency (Varma, Shah, Kenton, Kramár & Kumar, 2023)
- arXiv · Grokking at the Edge of Numerical Stability (Prieto, Barsbey, Mediano & Birdal, 2025)