Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets

To study generalization in a clean, controlled setting, the authors chose algorithmically generated datasets. These are not natural images or text — they are arithmetic problems with a known, exact answer. The core task: given two numbers $a$ and $b$, predict the result of a binary operation like $(a \circ b) \bmod p$, where $p = 97$ is a prime number.

Why modular arithmetic? Because the complete dataset is finite and small. For a prime $p = 97$, there are $97 \times 97 = 9{,}409$ possible input pairs. The training set is a random fraction of this table — say 50%. The test set is the rest. There is no ambiguity: every answer is deterministic.

The paper tested a variety of operations, including addition, subtraction, division, polynomial combinations like $x^2 + xy + y^2 \bmod p$, and even composition in the symmetric group $S_5$ (permutations of 5 elements). This diversity lets the authors check whether grokking is specific to one operation or a general phenomenon.

The model is a small decoder-only Transformer: 2 layers, 128-dimensional embeddings, 4 attention heads, and approximately 290,000 parameters. The input tokens represent the two operands $a$ and $b$, and the network predicts the result from 97 possible output classes.

This model is massively overparameterized for the task — 290K parameters to learn a function with fewer than 10K examples. In classical machine learning, this should guarantee overfitting and poor generalization. And indeed it does — at first. What happens next is the surprise.

Here is what the authors observed. When training on division modulo 97 with 50% of the data, the training accuracy reaches 100% within a few hundred epochs. At this point, the validation accuracy is barely above random chance — the network has simply memorized the training examples without learning any generalizable pattern.

If training stopped here, the conclusion would be textbook: the model overfit. But the authors kept training — for thousands, then tens of thousands of additional epochs. And then, something remarkable happened. After roughly $10^{4}$ to $10^{5}$ epochs of stagnation, the validation accuracy suddenly shoots up from near-chance to near-perfect. The network transitions from rote memorization to genuine generalization — and it does so abruptly, like a phase transition in physics.

The authors named this phenomenon grokking, after Robert Heinlein's term for deep, intuitive understanding. The network doesn't just learn the training examples; it groks the underlying mathematical structure.

Grokking training can be divided into three distinct phases, each with a different relationship between training and validation performance:

Phase 1 — Memorization. The network rapidly fits the training data. Training loss drops to near zero, training accuracy reaches 100%. But validation accuracy remains at chance level. The network has stored individual input-output mappings without discovering any underlying rule. Think of it as a lookup table — fast, perfect on known inputs, useless on anything new.

Phase 2 — Plateau. Training metrics stay perfect while validation shows no improvement. This phase can last for orders of magnitude more epochs than Phase 1. It appears that nothing is happening, but the internal representations are slowly reorganizing. Weight decay is gradually simplifying the network's parameters, pushing the weights toward solutions that encode patterns rather than individual examples.

Phase 3 — Generalization. Validation accuracy suddenly rises — often from near-chance to near-perfect in a relatively short window of epochs. The network has discovered the algebraic structure of the operation. For modular addition, follow-up research by Nanda et al. (2023) showed the network learns discrete Fourier components — a genuinely mathematical representation.

One of the paper's key findings is how dataset size affects the time to generalization. The relationship is far from linear: as the training fraction decreases, the number of epochs required to reach 99% validation accuracy increases dramatically — often by orders of magnitude.

At 70–90% of the data, training and validation curves track each other closely and generalization happens quickly. At 50%, there is a clear gap, and grokking takes thousands of extra epochs. Below about 30%, generalization may not happen at all within the optimization budget.

This suggests the existence of a critical dataset size — a threshold below which the network cannot discover the underlying pattern, no matter how long you train. Above this threshold, the network will eventually generalize, but the time required grows sharply as you approach it from above.

The paper found that weight decay is crucial for grokking. Without it, the network memorizes and stays stuck. With it, grokking happens reliably. But why?

Weight decay adds a penalty proportional to the squared magnitude of the weights. At each training step, the weights are pulled slightly toward zero. During the memorization phase, this has little visible effect — the network compensates by adjusting other parameters. But over thousands of epochs, this constant pressure gradually erodes the large, specialized weights that store individual examples and favors smaller, more structured weights that encode general patterns.

Think of it as a sculptor working on a block of marble. In the beginning, the raw shape (memorization) emerges quickly. But the sculptor keeps refining, chipping away excess stone (large weights) until the smooth, essential form (generalization) is revealed. Weight decay is the chisel.

Not all operations grok equally. The paper tested 12 binary operations and found that simpler, more symmetric operations grok more reliably:

Operations that grok consistently include $x + y$, $x - y$, $x / y$ (mod $p$), $x^2 + y^2$ (mod $p$), and group composition in $S_5$.

Some polynomial operations like $x^2 + xy + y^2$ (mod $p$) and $x^3 + xy$ (mod $p$) did not grok within the training budget. These operations have more complex algebraic structure that may require either more training time, more data, or different model configurations.

A key insight: operations with more inherent symmetry tend to generalize faster. Modular addition is fully commutative and associative — properties that create redundancy in the training data, effectively increasing the information per example. The symmetric group operations, while non-commutative, have rich structural regularity that the network can exploit.

Grokking sits at the intersection of several deep puzzles in deep learning.

Double descent. Classical bias-variance theory predicts that test error should rise with model complexity once you pass the interpolation threshold. But modern deep learning shows a second descent: very overparameterized models generalize better. Davies et al. (2023) showed that grokking can be viewed as epoch-wise double descent — the same phenomenon playing out over training time rather than model size. The model first reaches an interpolation peak (memorization), then descends again to low test error (grokking).

Emergent abilities. Large language models sometimes exhibit sudden capability jumps at certain scales — tasks that go from random-chance to high accuracy over a narrow range of model sizes. Grokking suggests that similar phase transitions can happen along the time axis rather than the scale axis. Both phenomena share the pattern: prolonged apparent failure, then sudden success.

Mechanistic interpretability. Nanda et al. (2023) used grokking as a case study for mechanistic interpretability. They showed that a grokked network performing modular addition learns to embed inputs as points on a circle and computes addition via rotation — a discrete Fourier transform implemented in the network's weights. This was one of the first demonstrations that neural networks can learn interpretable algorithms, not just statistical patterns.

Grokking is surprisingly easy to reproduce. Below is the essential setup: generate a modular addition dataset, split it, train a small Transformer with weight decay, and wait. The key ingredients are the small dataset, the overparameterized model, and sufficient weight decay.

Four years after its discovery, grokking has become a touchstone in deep learning theory. It demonstrates that neural networks are capable of far more nuanced learning dynamics than the simple "fit and overfit" narrative suggests. The phenomenon sits at a crossroads of generalization theory, mechanistic interpretability, and optimization dynamics — and continues to yield insights about the nature of learning itself.