grokking, from scratch, in the browser

a neural network memorizes, then suddenly understands

let's say you show a neural network half the answers to an addition table (a + b mod p). it memorizes them perfectly. train accuracy 100%, test stays near zero. looks like it just memorized, right?

then you keep training. for a long time nothing seems to happen. and then, thousands of steps later, test accuracy jumps from near-zero to near-perfect. the network figured out the rule, not just the answers. that's grokking.

why does it happen? the general rule turns out to be a handful of sine waves (Nanda et al. 2023), which is a cheaper way to fit the data than a lookup table of every pair. weight decay penalizes expensive solutions, so it tips the network toward the rule. the hidden progress chart tracks how much the network relies on those sine waves.

you can try an MLP or a Transformer (the original paper used a Transformer). both are written from scratch in javascript. the two knobs that matter most are weight decay and training fraction. try the presets below.
presets:
press to watch the network learn
Memorizing Generalizing Grokked!
0
training step
0
steps/sec
0.0%
train accuracy
0.0%
test accuracy

Every answer, a + b before training

correct wrong no prediction yet
training pair (dot) held out (no dot)
x-axis = a, y-axis = b. watch the held-out cells turn white.

Accuracy

grey = seen data, white = unseen data. watch the white line.

Loss

lower is better. watch test loss plateau then crash — that's the transition.

Hidden progress what's happening during the "flat" part

fourier concentration = share of the network's weight energy in its top 5 frequencies (random weights sit near 5 / (p/2)). when it rises, the network is building the sine-wave solution; compare its timing with test accuracy. weight norm = sum of squared weights. try toggling weight decay and watch how differently it grows. dashed lines: grey = memorized (train ≥ 99%) and white = grokked (test ≥ 90%). the gap between them is the grokking delay.

Fourier spectrum what frequencies the network uses

each bar = one frequency component (k=1 to p/2). after grokking, only a few bars should spike. those are the frequencies that solve modular addition.

Ask the network idle

+ mod 97
Network: not initialized
Try this before and after grokking. Early on, it memorizes training pairs but guesses randomly on new ones. After grokking, it gets everything right.

What the neurons learn first-layer weights on a clock

each circle = one hidden neuron. the p points around the clock show its weight for each input value. smooth waves = Fourier features (the network found the circular structure of mod p).
before (step 0)
now (step 0)

Learned representations PCA of weight vectors

each dot = one hidden neuron's weight vector, projected to 2D via PCA. exploratory: expect loose structure, not clean clusters. the clock plots above are the clearer view (switch to Transformer to see the famous embedding circle). color = index.
before (step 0)
now (step 0)