II · THE IDEA · ARTIFICIAL INTELLIGENCE
Mechanistic Interpretability of Grokking
▶ Listen · narrated
A model can score perfectly on its training set for thousands of steps and still fail every held-out example. What changes inside when the test accuracy finally rises is now partly visible.
At a glance
- Phenomenon
- Generalisation long after train accuracy has already saturated
- Classic task
- Modular addition on small algorithmic tables
- Key method
- Mechanistic reverse-engineering of learned circuits
- Reported phases
- Memorisation, circuit formation, then cleanup
- Progress tools
- Restricted and excluded loss beyond train/test curves
Think of a student drilling a times table. At first they memorise every sum they have seen and score full marks on those drills, yet freeze on any new pair. Much later, after far more practice, they suddenly handle unseen pairs too—not because they stored more answers, but because they acquired a method for adding.
Grokking is the version of that story seen in small neural networks on tasks such as modular addition. The training score hits perfect early. The test score stays near chance for a long time, then rises sharply while training continues. Mechanistic interpretability opens the network and asks what changed. In the well-studied modular-addition case, researchers found a structured method using frequency patterns and circle-like addition rules—not just a pile of memorised pairs. Training appears to pass through stages: memorise first, build the general method underneath, then clean away the pure memorisation so the method can drive answers. The late jump on the test curve is often that cleanup becoming visible, not the first moment the method exists.
Power et al. documented grokking on small algorithmic datasets: models reach near-zero training loss and perfect train accuracy while validation accuracy remains near chance, then—after extended further optimisation, often with weight decay—validation accuracy rises sharply. Nanda et al. reverse-engineered a one-layer transformer on modular addition and recovered a sparse Fourier / trigonometric algorithm implementing addition via phase structure (compositions equivalent to angle addition).
They introduced progress measures beyond train and test loss. Restricted loss evaluates the isolated generalising circuit; excluded loss evaluates the model with that circuit ablated. These measures reveal three phases in weight space: (1) memorisation—train loss drops, generalising circuit still weak; (2) circuit formation—the Fourier addition circuit strengthens while memorising components persist; (3) cleanup—memorising components diminish and test loss falls as the structured circuit dominates. The externally sudden generalisation jump is therefore compatible with a slower internal handoff between algorithms. Limitations are important: the detailed circuit and phase decomposition are tied to this task class and architecture; train/test curves alone do not identify phases; and regularisation such as weight decay is empirically linked to grokking in the source work without by itself specifying which circuit will win.
Look closer
Train and test curves come apart
On small modular arithmetic datasets, training accuracy can reach perfect scores while test accuracy remains near chance for a prolonged stretch of further optimisation. Only much later does test performance rise sharply toward the training level. The delay is the defining empirical signature of grokking: overtraining continues after memorisation is complete, and generalisation arrives as a later event rather than as a smooth companion of falling train loss.
A Fourier-style addition circuit
In a one-layer transformer trained on modular addition, reverse-engineering work recovered an algorithm that embeds inputs in sparse Fourier features and composes them with trigonometric identities equivalent to addition on a circle. The network is not looking up arbitrary memorised pairs at the end of training; it implements a compact, structured procedure whose intermediate activations and weight structure match that procedure closely enough to be read out.
Three phases in weight space
Progress measures that isolate the generalising circuit from memorising components show a sequence rather than a single jump. Early on, the network memorises training pairs. A generalising circuit then forms while memorisation is still present. Cleanup follows: components that served only memorisation shrink, and test loss falls as the cleaner algorithm dominates. Train and test loss alone blur these stages; the restricted and excluded losses make them visible.
The story
Grokking entered the literature as an empirical surprise on small algorithmic datasets. Networks trained to predict the result of binary operations—modular addition among them—could fit every training example long before they performed well on held-out pairs drawn from the same operation. Optimisation did not stop when train loss neared zero. After a further stretch of training, often under weight decay, test accuracy rose abruptly from near chance toward perfect generalisation. The shape of the curves, not merely the final performance, was the point: generalisation arrived late, and looked discontinuous on ordinary plots.
That late jump invited a mechanistic question. If the network had already memorised the training table, what later change in the weights produced generalisation rather than more memorisation? Work on progress measures for grokking attacked this by reverse-engineering a one-layer transformer on modular addition. The recovered algorithm did not resemble a lookup table. Inputs were represented with sparse frequency components; addition was implemented through phase structure and trigonometric identities that amount to adding angles on a circle. Once identified, that circuit could be measured directly during training instead of being inferred only from the test curve.
The measurements split training into phases that ordinary loss curves smear together. In the memorisation phase, train loss collapses while the generalising circuit is still weak, so test performance stays poor. During circuit formation the Fourier-style procedure strengthens even though memorising components remain. Cleanup is the stage in which those memorising components are removed or suppressed; test loss then falls as the network comes to rely on the structured algorithm. Restricted loss tracks how well the isolated generalising circuit performs; excluded loss tracks performance when that circuit is ablated. Together they show that the dramatic test jump is often the visible tip of a slower internal transition already under way in weight space.
The editorial force of this line of work is modest but sharp. Grokking is not explained here as a universal law of deep learning. It is explained, for a concrete task and architecture, as competition and succession between algorithms that leave different traces in the weights. Memorisation can win early. A generalising circuit can form underneath it. Cleanup can hand control to that circuit long after train accuracy has nowhere left to go. Seeing those phases requires tools that look inside the network, not only at the two scalar curves that first made the phenomenon famous.
Weight decay and related regularisation matter in the original reports because they favour the cleaner solution over pure memorisation, but the mechanistic account does not reduce to “regularisation causes generalisation.” It specifies what is being favoured: a particular Fourier addition algorithm rather than an opaque store of pairs. That specificity is what makes the result usable as interpretability rather than as another training anecdote. Once the circuit is known, one can ask whether a given checkpoint has built it, how strong it is relative to memorising pathways, and which training choices accelerate formation or cleanup. Those are questions about structure in weight space, answered with progress measures tied to a reverse-engineered mechanism.
Why it mattered then
When grokking was reported, the dominant habit was to read generalisation from the gap between train and test loss. A network that had already driven train loss to the floor and only later lifted test accuracy did not fit tidy stories in which generalisation tracks continuous compression or early stopping. The phenomenon mattered in its moment because it forced a separation between “fit the training set” and “acquire a procedure that extends.” Mechanistic follow-up mattered for a second reason: it showed that the late jump could be decomposed into identifiable algorithmic stages rather than left as an opaque phase transition. For interpretability research, that was evidence that reverse-engineering could turn a striking training curve into a sequence of concrete circuits and measurable progress, even when the external behaviour looked sudden.
Why it matters now
Sudden shifts in capability during extended training remain a live concern whenever models are trained far past the point of low training loss. Grokking on modular addition is a narrow setting, yet it supplies a worked example of how to refuse a purely external reading of such shifts. If a jump in test performance can hide a memorisation phase, a circuit-formation phase and a cleanup phase, then capability curves alone are incomplete instruments. The same discipline—name the algorithm, build progress measures that track it, watch components compete in weight space—applies whenever one wants to know whether a model has formed a general procedure or merely stored training patterns. The modular-addition case remains useful less as a claim about all networks than as a template for that discipline.
The surprising detail
The generalising solution recovered for modular addition is not a softened lookup table. It is closer to a clock: sparse Fourier features and trigonometric composition that implement addition as rotation. For a stretch of training this structured circuit coexists with memorisation; test accuracy stays poor until cleanup removes the memorising pathway. The dramatic generalisation jump can therefore lag the birth of the algorithm that eventually explains it.
What is disputed
The phased circuit account is established in detail for small transformers on modular addition and related algorithmic setups. How far the same memorisation–formation–cleanup sequence describes grokking in larger models or non-algorithmic tasks remains open; the original grokking reports and the mechanistic progress-measure work should not be flattened into a universal training law.
Remember this
Grokking’s late test jump can mark cleanup after a generalising circuit has already formed beneath memorisation—not the instant the algorithm is born.
Test yourself
A modular-addition network has perfect train accuracy and still-poor test accuracy. Restricted loss for the Fourier addition circuit is already low, while excluded loss remains bad. Which phase does this pattern suggest, and what would you expect to happen to test accuracy if cleanup proceeds?
Circuit formation (or the cusp before cleanup): the generalising circuit already works when isolated, but memorising components still dominate the full model, so test accuracy stays low. If cleanup proceeds and those memorising components shrink, the full model should come to rely on the circuit and test accuracy should rise sharply—even though train accuracy was already perfect and the circuit itself may change little.
Go deeper
- [2201.02177] Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets · arxiv.org
- [2301.05217] Progress measures for grokking via mechanistic interpretability · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.