Capstone: grokking & where to go next

The puzzle: would you have stopped this run?

Imagine you are babysitting a training run. A model trains on modular addition: the week's specimen is a one-layer transformer fed three tokens [x, y, =] and trained to predict (x + y) mod 113 (113 is prime, which keeps the math clean). It sees only a fraction of the possible (x, y) pairs during training; the rest are held out as a test set.

Training accuracy soon reaches 100%, so the model gets every observed pair right. Test accuracy remains near chance and reveals no transfer to held-out pairs. The dashboard looks like classic overfitting: a memorized training table, flat curves, and thousands of steps without visible progress. You would probably stop the run and tune the setup.

Continuing the same run reveals Grokking . Thousands of steps after training accuracy saturates, test accuracy climbs sharply toward 100%. The optimizer, data, and learning-rate schedule remain unchanged while the model shifts from memorized cases toward a general algorithm. Power et al. first reported this delayed generalization in 2022. This week you compare the model before, during, and after that transition.

Nanda et al. explain the delay through two competing solutions inside one network. A memorized lookup table forms quickly and uses substantial weight to store arbitrary cases. A compact general algorithm develops slowly. Weight decay continuously shrinks parameters, so the cheaper algorithm gains an advantage as its circuit improves. Once it supports the task, continued regularization removes the costly memorized solution. The abrupt test-loss drop marks the visible end of a gradual handover. The paper's progress measures expose that hidden work by tracking quantities in the weights that change smoothly throughout the plateau.

Two progress measures separate memorization from the Fourier algorithm. One measures how much model weight lies in the key Fourier frequencies used by the general circuit. Another tracks weight devoted to memorizing individual training examples. Across checkpoints, the Fourier component grows while the memorization component eventually shrinks. Plotting those measures against training step reveals motion during the flat test-loss plateau. Behavioral loss shows when the handover finishes; weight-space measures show it being built.

A training analysis needs checkpoints from the plateau as well as the endpoints. Keep the train/test split fixed and save model weights on a logarithmic schedule, with denser checkpoints near the expected transition. At each checkpoint, record full-train and full-test loss, accuracy, total weight norm, and the Fourier progress measures. Repeat the run across seeds because the grokking step can move substantially. A control run with weaker or zero weight decay tests the proposed selection pressure: memorization can remain competitive much longer when shrinking weights carries less reward.

The remaining question is concrete: which compact algorithm did the model learn? One piece of Fourier math supplies the answer.

Fourier for programmers

Modular arithmetic wraps around, so periodic waves fit its structure. Under mod 113, 112 + 1 returns to 0 like a clock hand passing twelve. The Fourier basis represents a repeating signal as a sum of sine and cosine waves. Here, a frequency counts complete cycles across the 113 possible values. Frequency 5 repeats five times as x moves from 0 to 112. The FFT (fast Fourier transform) reports the strength of each frequency in a signal.

You can verify that in four lines. Hide two frequencies in a signal and ask the FFT to find them:

# 01: the FFT finds hidden frequencies
import numpy as np

t = np.arange(113)
signal = np.cos(2*np.pi*5*t/113) + 0.5*np.cos(2*np.pi*17*t/113)
loud = np.abs(np.fft.rfft(signal)).argsort()[-2:]
print(sorted(loud.tolist()))
# [5, 17]

A periodic table can sometimes be described compactly by a small set of Fourier frequencies. The FFT identifies those frequencies.

To inspect an embedding matrix, treat each embedding dimension as a signal over token values 0 through 112. Apply the FFT along that token axis, square the magnitudes, and sum across embedding dimensions. A memorized representation spreads power across many frequencies. The grokked checkpoint shows sharp peaks at a small key set. Sine and cosine at the same frequency appear as a two-dimensional subspace, so the analysis measures their combined power and stays invariant to arbitrary phase or basis orientation.

Waves also turn modular addition into rotation. Represent x as a point on the unit circle at angle 2πx/113. Multiplying two such complex points adds their angles, and the circle handles wrap-around:

# 02: modular addition as rotation
import numpy as np

p = 113
def as_angle(x):
    return np.exp(2j * np.pi * x / p)   # x as a point on the unit circle

x, y = 100, 50
combined = as_angle(x) * as_angle(y)    # multiplying = adding the angles
angle = np.angle(combined) / (2 * np.pi) * p
print(round(angle % p))
# 37  # rotation computes (100 + 50) mod 113

The circle handles wrap-around geometrically, so the computation needs neither a conditional nor an explicit modulo operation. Sine and cosine give the coordinates of points on that circle, which lets a network express modular arithmetic through periodic components.

The representation uses both coordinates because angle addition needs phase. For angles a and b, the identities cos(a+b) = cos(a)cos(b) − sin(a)sin(b) and sin(a+b) = sin(a)cos(b) + cos(a)sin(b) turn products of input waves into a wave for their sum. The network approximates these products through attention and its MLP. The unembedding compares the resulting phase with each candidate token's phase, translating the internal rotation back into output logits.

The learned algorithm

The trained transformer embeds each number with sine and cosine waves at a few key frequencies. Attention and ReLU help multiply the waves for x and y, while trigonometric identities produce output-logit terms such as cos(w(x + y - z)) for candidate answer z. These terms align when z = (x + y) mod p and tend to cancel for other candidates, making the correct logit largest.

For each key frequency, cos(w(x + y − z)) reaches its maximum value of 1 when z ≡ x + y. At other candidates, phases from different frequencies spread around the circle and partially cancel. Summing several key frequencies gives the correct candidate the largest score:

# 03: constructive interference picks the answer
import numpy as np

p = 113
x, y = 41, 87
key_freqs = [14, 35, 41, 42, 52]        # a small set, like the ones the model learns
z = np.arange(p)
logits = sum(np.cos(2*np.pi*k*(x + y - z)/p) for k in key_freqs)
print(int(logits.argmax()), (x + y) % p)
# 15 15  # predicted candidate and exact modular sum agree

Gradient descent selected this compact trigonometric representation without receiving it in the objective. Under weight decay, the model encoded numbers as rotations, multiplied waves, applied identities, and scored answers through interference. The result extends week 7's geometric lesson from feature packing to an arithmetic algorithm. Mechanistic interpretability can state the representation, intermediate computation, and output rule in detail.

Each evidence source constrains a different part of that account. Embedding and unembedding spectra identify which frequencies enter and leave the network. Activation analysis checks whether hidden states contain products at sums and differences of those frequencies. Trigonometric identities predict which terms should contribute to a candidate logit. Ablations then test whether preserving those terms preserves accuracy. Agreement across weights, activations, algebra, and interventions makes the explanation specific enough to fail if any claimed step is wrong.

The reverse-engineering claim needs evidence from several views of the model:

• inspect weights and embeddings;
• analyze activations and frequency spectra;
• use Restricted ablation to remove everything except the hypothesized mechanism;
• compare behavior before and after interventions.

Restricted ablation gives the strictest test. Project the model's computation onto the hypothesized key frequencies, removing other Fourier components, and check whether task performance survives. Strong retained accuracy shows that the selected frequencies carry most of the algorithm. Weight inspection from week 2, activation analysis from week 6, and causal intervention from week 9 converge on the same mechanism.

Run complementary ablations for a sharper result. Keeping only key frequencies tests sufficiency. Removing those frequencies while preserving the rest tests necessity and should destroy modular-addition accuracy. A random set of the same size controls for the amount of retained computation. Report both task accuracy and the fraction of weight or activation norm preserved, since a frequency set that keeps nearly everything provides a weak restriction. The paper's small key set retains the mechanism despite discarding most alternatives.

Where to go next

Several skipped ARENA notebooks extend the same tools:

• 1.5.3 OthelloGPT: test whether a model trained on Othello move sequences builds an internal world model of the board, using week 10's probing skill on a new domain.
• 1.4.2 SAE Circuits: combine week 8's features with week 9's circuit methods by treating SAE latents as circuit nodes.
• 1.5.1 Balanced Brackets: reverse-engineer a small classifier the way you just reverse-engineered modular addition.
• 1.3.4 Activation Oracles: a modern take on reading model internals, if weeks 10-11 were your favorites.
• ARENA chapters 2-4 are "season 2" tracks: reinforcement learning, LLM evals, and alignment science.

Read Nanda et al.'s Progress measures for grokking via mechanistic interpretability, the original grokking paper by Power et al., and Neel Nanda's grokking analysis post, whose plots the notebook reproduces. Neel Nanda's getting-started guide links concrete open problems, and the MATS program offers a route into mentored research.

A good first original project is small: pick one behavior, one model, one metric, and one intervention. Reuse the tools you already touched: TransformerLens caches, SAE dashboards, probes, patching, or steering.

Pair-session guide

Core work is ARENA 1.5.2 sections 1 and 2: Periodicity & Fourier Basis, then Circuit and Feature Analysis. ARENA provides a trained checkpoint for analysis, so the session skips the long training run. Stretch work is section 3, Analysis During Training, which replays saved checkpoints to show the grokking dynamics. Leave time for a course retrospective: each person names one technique they can use independently and one next project they might attempt.

For the checkpoint stretch, use one consistent set of key frequencies derived from the final grokked model, then measure their strength at every earlier checkpoint. Recomputing a new favored set at each step can hide whether the final circuit existed earlier. Plot restricted-ablation accuracy alongside test accuracy and Fourier power. If the mechanistic story holds, final-circuit power and restricted accuracy should rise through the plateau before the unmodified model's test curve completes its sharp transition.

What you should see

You should see about half a dozen sharp spikes in the embedding matrix's Fourier spectrum, corresponding to the provided model's key frequencies. In the section 3 stretch, replayed checkpoints should show training performance saturating long before test performance improves.