Grokking

Have you ever had this experience when learning something new where you just couldn’t understand what was going on, so you just memorized all the answers to the questions that you think will show up on the test? But you were bored, so you kept memorizing and answering different variations of the test questions repeatedly even after you were pretty sure you memorized everything. And then one day, it occurred to you in a sudden flash of realization: Hold on a sec! There’s a pattern to this! I just have to understand this one idea or principle and if I apply it in this specific way, then I can get the right answer every time without memorizing everything!

If you have had this experience before, then you’ll understand what Alethea Power et. al. was trying to find out in their paper Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets¹. The authors asked the question: could this sudden flash or realization happen to neural networks? And their answer is that it does! Neural networks can learn to find this “general principle” that works to answer all the test questions without memorizing everything. The authors call this phenomenon “grokking”.

What Do You See When A Model Groks?

The most dramatic thing you’ll observe is the model’s accuracy on the training and validation sets. At first, it looks like the classic overfitting curve. The model first reaches almost 100% accuracy on the training set and saturates there, while validation accuracy remains no better than random chance. But if you let the model continue training for much longer, something interesting starts to happen. The validation accuracy also starts to increase. And if you let the training process run for long enough, eventually, the validation accuracy also reaches close to 100%.

Maybe this sounds like common knowledge now but it really wasn’t obvious at all that this would happen back when the paper was published. At that time, the machine learning community at large had the intuition that once a model had memorized its training set and validation accuracy was still struggling, continuing to train the model was unlikely to make it generalize any further. ML researchers usually treated this as a sign of overfitting (bad because we’re wasting compute on a model that’s bigger than it needs to be). The standard recommendation at the time would be to get more data, use stronger regularization, dropout, early stopping or a smaller model. The bitter lesson of “keep scaling compute until the model starts to get better” wasn’t such a dominant view at the time. In fact, the legend goes that the authors discovered grokking completely by accident after unintentionally leaving a model training over their holidays. Most researchers at the time probably wouldn’t intentionally wasted compute on this experiment. This brings us to some key elements of the experimental setup that made the original grokking curve so dramatic.

How Did They Set Up the Experiment?

The Dataset is Intentionally Finite and Generateable. This makes the experiement much easier because we know in advance the full dataset and we can easily generate it without spending a lot of effort going out into the real world to collect data. And since we know the rule to generate the full dataset, we can easily verify the model’s answer and that it is generalizing. It’s like showing someone half of a picture of an object like a computer mouse or a water bottle and asking them to predict what the whole thing looks like. In the case of the original grokking paper, the authors use various modular binary operations like addition, subtraction, division. And since the operations are modulo a prime number, we can generate the full finite table of the binary operations easily in memory.

The Model is Heavily Over-parameterized. This means that the model must have way more parameters than needed for the task so that it can achieve memorization easily. In the original grokking paper, the authors use a small decoder only Transformer with two layers, a width of 128 four attention heads and roughly 400,000 parameters for the simple task of predicting the result of a modular binary operation on two numbers. The now famous grokking curve was produced by the model trying to predict division of two numbers modulo 97. The over-parameterization doesn’t seem to be requirement for grokking though. Research that came after seems to suggest that grokking is more dependent on the precise nature of the learning dynamics and the interplay between different solution types within the model’s parameters. Grokking occurs in simple learning models like logistic regression where you essentially only have one layer in your neural net.²

The numbers are not fed into the model as numbers. Instead, they are tokenized and the model sees the inputs as a series of seemingly meaningless tokens. This is a deliberate and unexpected choice by the authors. The model does not have two integer or floating point inputs and one answer output. Instead, every residue modulo 97 is represented by its own independent token. And as far as the model knows, the token representing 13 is no more closely related to 18 than it is to 70. The authors explicitly describe the elements of the operation as discrete symbols with no internal structure. What the model actually sees is a sequence of tokens roughly of the form:

<x> <operation> <y> <equals> <answer>

This is the equivalent of handing someone a multiplication table written in an alien language with a missing row and then asking them to fill it in.

You have to keep training it for several orders of magnitude of steps after memorization. The model reaches perfect training accuracy after around 1000 steps. But the grokking phenomenon starts to become apparent only after 100,000 steps or so of training, and in the paper’s experiment the model reaches high validation only after about a million steps. And it could take even longer. When I was trying to replicate the experiment, I found that even changing small things like the way the weights are initialized or the order in which the RNG is consumed can impact the grokking curve. One of my runs needed more than 10,000,000 steps to show any signs of grokking happening.

They intentionally hold back some well known tricks to make models generalize. For the most dramatic experiment the authors actually used the Adam optimizer with no weight decay rather than their usual AdamW optimizer to train the model. The authors also mention that learning-rate annealing improved performance in some experiments, but they deliberately chose not to use it in the main experiments, partly for simplicity.

Can You Grok More Complicated Things?

A natural next question might be: if over-parameterized models can grok binary operations, can they grok more complicated things? The original grokking paper answers this question partially. They test a variety of binary operations as well as several polynomial functions to see how grokking happens. Some modular binary operations have properties that make them effectively “equivalent” from the neural network’s perspective. For example, addition modulo (p-1) and multiplication modulo (p) are equivalent up to a relabeling of the symbols, and the same is true for subtraction and division. The authors find that these equivalent operations behave similarly when it comes to how much data they need to generalize. Some operations where the order of the operands doesn’t matter—addition and multiplication, for example—also tend to need less data to generalize than closely related operations where the order does matter, like subtraction and division.

But for some of the more complicated operations, the model struggled much more. In particular, for the polynomial x^3 + xy^2 + y mod 97, the authors could not get the model to generalize within their allowed optimization budget at any training-set size up to 95% of the full dataset. The model effectively just memorized the examples it had seen without finding a pattern that generalized to the remaining ones. So it does seem that, at least with the model and training setup used in this paper, some more complicated and less symmetric functions are much harder to grok.

There are also some interesting evidence in Towards understanding grokking: An effective theory of representation learning from Ziming et. al. where the authors seem to achieve grokking on the MNIST dataset for handwritten digit recognition. This is a non-algorithmic dataset. I’ll talk a little more about their paper later on in this article. So grokking seems to be easier to achieve on certain types of complexity.

What’s Really Happening During Grokking?

So what’s really happening during grokking? Some of the follow-up papers published since then have tried to give us an understanding. In the paper Progress Measures for Grokking via Mechanistic Interpretability³, Neel Nanda and his colleagues tried to understand what exactly is happening inside the neural network as the model achieves grokking. What they essentially find is that the Transformer learns an actual algorithm for performing modular addition. And surprisingly, that algorithm looks a lot like Fourier analysis and trigonometry. In the paper, the model is trained on addition of two numbers a and b modulo 113, and the model seems to learn to produce sine and cosine components of a and b at a bunch of different frequencies.

$$\left[\cos\left(\frac{2\pi k}{113} a\right),\sin\left(\frac{2\pi k}{113} a\right)\right], \qquad \left[\cos\left(\frac{2\pi k}{113} b\right),\sin\left(\frac{2\pi k}{113} b\right)\right]$$

So the network tries to represent the modular addition of these numbers using trig identities like

$$\cos\left(\frac{2\pi k}{113} a\right) \cdot \cos\left(\frac{2\pi k}{113} a\right)$$

So instead of memorizing (37+52=89) the model has discovered a representation in which addition becomes a little easier. It then compares this representation of (a+b) against the sine and cosine representation of each possible answer (c). Roughly speaking, it produces terms like:

$$\cos\left(\frac{2\pi k}{113}(a+b-c)\right)$$

which become largest when c=a+b mod 113. With several frequencies combined together, the model can identify the correct answer very precisely. Nanda and his colleagues verify this by doing Fourier analysis on the model’s weights and activations and by selectively removing Fourier components from the model to see which ones actually matter to the model’s predictions.

This is one of the coolest parts of grokking to me. Nobody explicitly set up the model architecture to model sine waves or use trig identities. The training removed as much structure from the inputs as possible but it was still able to produce or “discover” these relationships, building them out of pure chaos, nudging the model weights just so here and there to take advantage of the structure of mathematical objects. Welch Labs has a great video on YouTube explaining this visually, including some beautiful animations and a great explanation of our best intuition behind what transformers learn when grokking.

Another interesting thing Nanda and his colleagues talk about in the paper is that just looking at the training and validation accuracy curves don’t tell the full story. The way the curves evolve as the model trains leads you to think that the model isn’t doing anything interesting between memorization and generalization and that it happens suddenly. he paper’s analysis, however, suggests that the model does this in several steps: memorization, circuit formation, and then cleanup. Basically, the model first memorizes, and while it seems like nothing is happening afterward, it is forming structures called ‘circuits’ that perform the more general algorithm for modular operations. And then finally, there’s a “cleanup” phase where the model removes the old parts that memorized the answer.

There’s another paper that was published not long after the Grokking paper that tries to understand a little deeper about what causes grokking behaviour. In their paper Towards Understanding Grokking: An Effective Theory of Representation Learning⁴, Ziming Liu and his colleagues approach grokking from a physicist’s point of view.

They discovered that when the model groks, its internal representations are way more structured than in models that simply memorize, and this structure seems to reflect some of the underlying structure of the problem that the model is being trained on. For example, in their Figure 1 they project the embedding that the model learned down to a dimension of 2 using PCA at different stages. At memorization, the components look essentially randomly distributed but after the model has generalized, the embeddings have organized themselves into a very clear circular structure that reflects the structure of modular addition.

Even more interesting is the second part of the paper that treats the state of the weights of the network is something of a physical material like water that undergoes phase transitions under different physical conditions. The authors identify four different phases: comprehension, grokking, memorization and confusion. In the comprehension phase, the model learns to generalize about as quickly as it learns to memorize. In the grokking phase, the familiar delay between memorization and generalization from Power et. al. appears. In both comprehension and grokking phases, the model is able to generalize. In the memorization phase, the model fits the training data well but never reaches good validation performance. And in the confusion phase, the model cannot even reliably memorize the training set. The paper formalizes these definitions using when training and validation accuracy cross 90%; “comprehension” means that those two events occur within 1,000 training steps of each other, while “grokking” means that both happen but with a longer delay.

These phases don’t appear to be scattered randomly across hyper-parameter space either. They form very clear regions and in the authors’ toy models, grokking tends to sit between comprehension and memorization. This analysis makes a lot of sense because by my understanding, the experiments in the original grokking paper were set up carefully to emphasize the dramatic delay between memorization and generalization.

Grokking May Be A General Feature of Learning Systems

So what do we make of all this? Grokking as a phenomenon feels eerily human to me. You first have the model memorize answers to get an easy reward from the optimizer, then as we train / punish the model more, it works hard to find a general principle. But when we looked inside the model, we found that the generalization doesn’t really appear out of nowhere in a flash of “inspiration”. Rather, it seems like the initial training run produces some seeds of an idea for a more general solution that gets refined further as you let the model train.

Some recent research also seems to show that it may not even be necessary for a system to be a neural network or a transformer to show grokking behaviour⁵. There’s also some indications that grokking can emerge under very different kinds of regularization and inductive biases⁶. So the phenomenon doesn’t seem to be limited by architecture and learning method. It just seems to be a fact that when you try to learn a pattern, our model — whatever it is — first finds the easy solution then as you keep learning, finds a more general one. I wonder if grokking is a general feature of anything that learns — including neurons across life on Earth — and whether some similar phenomenon governs our own learning processes.

So coming back to cramming for tests, maybe my analogy was a little wrong. The student trying to practice for a test doesn’t suddenly discover the pattern in a magical instant. As they practice, maybe their mind — consciously or subconsciously — starts piecing together parts of the general pattern. They understand it piece by piece, chunk by chunk until the “aha!” moment appears. To the outside observer, and maybe even the student themselves, this may look like a flash of inspiration but inside their mind, their neurons were working hard to find a better way to learn what you were practicing.

References

Footnotes

1. Power, Alethea, et al. “Grokking: Generalization beyond overfitting on small algorithmic datasets.” arXiv preprint arXiv:2201.02177 (2022). ↩

2. Beck, Alon, Noam Levi, and Yohai Bar-Sinai. “Grokking at the edge of linear separability.” arXiv preprint arXiv:2410.04489 (2024). ↩

3. Nanda, Neel, et al. “Progress measures for grokking via mechanistic interpretability.” arXiv preprint arXiv:2301.05217 (2023). ↩

4. Liu, Ziming, et al. “Towards understanding grokking: An effective theory of representation learning.” Advances in Neural Information Processing Systems 35 (2022): 34651-34663. ↩

5. Mallinar, Neil, et al. “Emergence in non-neural models: grokking modular arithmetic via average gradient outer product, 2024.” URL https://arxiv.org/abs/2407.20199. ↩

6. Notsawo, Pascal Jr Tikeng, Guillaume Dumas, and Guillaume Rabusseau. “Grokking beyond the Euclidean norm of model parameters.” arXiv preprint arXiv:2506.05718 (2025). ↩