Grokking on modular addition

Grokking is delayed generalization: the model first memorizes—train accuracy shoots up and train loss collapses—while still doing poorly on held-out examples. Only after a long stretch of training does it generalize, often with a sharp rise in test accuracy on unseen cases. I trained a small decoder-only transformer on modular addition mod 97 (all 9,409 pairs; 40% used for training) with AdamW and weight decay 0.3.

The plot shows the usual grokking shape. Train accuracy rises almost immediately; held-out accuracy stays near zero for tens of thousands of steps, then climbs to 100% on the full unseen table. Loss matches that: train loss drops early while held-out loss stays high until the same late transition—a long gap between memorizing the training pairs and generalizing to the rest.

Code and notes are in github.com/fjavid/psi.