A one-layer transformer with four attention heads is trained to answer a + b mod p, seeing only a random 30 percent of all input pairs. It receives three tokens, a, b and =, and reads its answer off the last position. Training uses full-batch AdamW with a large weight decay.
The training curve has two stages. Within a few hundred epochs the model fits every training pair, but its accuracy on held-out pairs stays near chance, so it has memorized rather than learned the rule. Then, after a long plateau, test accuracy climbs from near zero to near perfect over a short window. That delayed jump is called grokking. Weight decay is the usual explanation: the memorizing solution needs large, unstructured weights, and decay slowly trades it for a smaller solution that actually computes the sum.
The smaller solution is periodic. Instead of a lookup table, the model represents each number as a point on a few circles, one per frequency k: the embedding of a contains cos(2πka/p) and sin(2πka/p). Adding angles on a circle is the same as adding numbers mod p, so the sum can be read out with trigonometry. The Fourier panel makes this visible: before grokking the spectrum of the token embedding is flat, and afterwards it concentrates on a handful of frequencies. The attention heads specialize in the same way, and the amount of attention a head pays to a becomes a periodic function of a.
The slider moves over saved checkpoints, which are denser where accuracy changed quickly. The loss and accuracy curves are recorded every ten epochs during training; the vertical line marks the current checkpoint, and clicking on a curve jumps to the nearest one. Each attention panel shows one head: the average pattern over all pairs (rows are query positions, columns are the positions attended to, only the last row matters for the answer), the pattern of the = query for eight fixed held-out examples, and the average attention to a as a function of a. The spectrum bars give the norm of the embedding at each Fourier frequency, normalized so their squares sum to one; the unembedding and the MLP neurons (as functions of both inputs, constant term removed) are available with the toggle. The weight norm tile is the L2 norm of all parameters, which weight decay pulls down during the plateau.
Reproduce: python3 train.py then python3 analyze.py, which rebuilds this page with the new run inlined. There is no export to GIF; step through frames with the arrow keys and take screenshots, or let it play.