Grokking, Visualized
Chapters
  1. The whole picture, first
  2. Counting on a clock
  3. The task as a dataset
  4. Vectors and matrices
  5. From scores to probabilities
  6. The model, step by step
  7. Simplifying the model
  8. How the model learns
  9. Grokking
  10. Angles and waves
  11. The Fourier basis
  12. Measuring fit and sparsity
  13. The Fourier multiplication algorithm
  14. Evidence 1: everything is periodic
  15. Evidence 2: neurons and heads
  16. Evidence 3: reading off the answer
  17. Evidence 4: interventions
  18. Progress measures and the three phases
  19. Watch it form
  20. Gradients and weight decay
  21. Is it robust? Seeds and controls
  22. The whole picture, again

Chapter 13 of 0–21

Evidence 1: everything is periodic

The weights and activations of the trained network, seen through the Fourier basis.

Sources for this chapter

Mechanistic interpretability tries to reverse-engineer the algorithm inside a trained network from its weights and activations. The paper supports its claim with four lines of evidence [1]; this chapter is the first: everything in the trained network is periodic, and sparse in the Fourier basis. All figures show the paper's network at the selected step (default: step 40,001, the model the paper analyses) — use the dashboard to move through training.

The embedding uses only a few frequencies

Each of the 128 rows of W_E is a list of 113 numbers — one per input number. Taking the Fourier transform of every row and measuring, for each frequency, the total size across the 128 rows gives the embedding's spectrum:

The embedding in the Fourier basis

Length of each Fourier component of the embedding matrix W_E at the selected step (blue cos k, red sin k).

Loading…

How to read it · the math
\big\lVert (W_E F^\top)_{:,i}\big\rVert_2,\quad i = \text{Const}, \cos 1, \sin 1, \ldots, \sin 56

W_E maps each number to a 128-dimensional vector. Reading those vectors as 113 functions of the number and taking their discrete Fourier transform (DFT) tells us which waves the embedding contains. Click a bar to select that frequency.

What to expect: After training only a handful of frequencies have non-zero bars: the key frequencies.

Learn step by step → · Source: [1] Fig. 3 (left); [3] “Understanding the embedding” · Computed by metrics.we_fourier_norms (code)

Played through training, the same norms become a race: the twelve largest at each moment, each with an outline at its starting value. The key frequencies pull ahead early, and the others shrink during cleanup.

The frequency race

The 12 largest frequency norms of W_E (or W_L) at the selected step, sorted, each with its value at initialisation as an outline.

Loading…

How to read it
<div class="fig-read"><p>Each bar is one frequency k: the size of the matrix's component along the cosine and sine of that frequency. Press ▶ to watch the bars reorder through training. Dark labels mark the run's key frequencies (chosen on the final model).</p>

What to expect: At initialisation all frequencies are about equal; the key frequencies pull ahead early in training, and the others shrink during cleanup.

Learn step by step → · Source: [1] Fig. 3 · Computed by checkpoint_metrics.py (code)

At the end of training only six frequencies carry visible weight in the paper's network: 14, 31, 35, 41, 42, 52. Five of them — 14, 35, 41, 42, 52 — reappear in every later part of the network; the paper calls them the key frequencies:

“We dub these the key frequencies of the model.”— Nanda et al. 2023, Sec. 4.1, p. 4 · source

(The sixth, k = 31, shows up in the embedding but not in the later layers; the paper notes it without naming it, and ARENA's text lists it [3].) Out of 56 possible frequencies the network uses 5 — and which five is arbitrary: other random seeds pick different ones (chapter 20).

What does a key frequency mean for the numbers themselves? Project every number's embedding onto the plane that carries cos(w_k a) and sin(w_k a): the 113 numbers lie on a circle, in the order k·a mod 113. Hover a number to find it on every circle.

Numbers on circles: the embedding in each key plane

Every number's embedding projected onto the 2D plane that carries cos(w_k a) and sin(w_k a), one panel per key frequency.

Loading…

How to read it
<div class="fig-read"><p>Each number n is coloured by the clock position the algorithm predicts for it at frequency k, namely k·n mod 113. A ring whose colours run smoothly around it means the embedding places n at angle w_k·n — the clock picture. Hands point at a, b and a + b; hovering a number enlarges it on every clock. While playing, the dots move between stored frames along the circle (polar blending); paused, the picture is a stored frame.</p>

What to expect: A cloud early in training that becomes a clean circle during circuit formation.

Learn step by step → · Source: [1] Sec. 3.1; [10] “Linear Probes” · Computed by geometry.key_plane_coords (code)

Control

A random control. Is this sparsity special? A control (experiment) answers that: compute the same spectrum for the network at initialisation (move the slider to step 1 in the dashboard) — every frequency has similar weight. The Gini coefficient of the spectrum is 0.04 at initialisation and 0.74 at step 40,001.

Attention patterns are periodic

For each head, the fraction of attention “=” pays to a, for every input (a, b), forms a striped 113 × 113 picture:

All four heads

Attention from '=' to a for every head, side by side (click one to select it).

Loading…

How to read it
<p class="fig-meta">
  <a href="../chapters/13-evidence-periodicity.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 9
  · Computed by <code>browser (GrokModel)</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Attention to a, for every input

Fraction of the selected head's attention that '=' pays to the first number a, for all (a, b).

Loading…

How to read it · the math
A^j_0(a,b) = \sigma\big(C_j[a] - C_j[b]\big)\ \text{(attention to '=' itself is negligible)}
<p class="fig-expect"><strong>What to expect:</strong> A periodic stripe pattern at the head's frequency; a + b symmetry.</p>
<p class="fig-meta">
  <a href="../chapters/13-evidence-periodicity.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 4 (left), Fig. 9
  · Computed by <code>browser (GrokModel)</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The paper reports head 0's pattern to be periodic with frequency 35; our analysis gives head frequencies 35, 42, 52, 42 for heads 0–3 (paper Table 2: 35, 42, 52, 42). Hover the grid: the value depends on a and b through waves of one frequency.

Neurons are periodic

The activation of a neuron over all inputs is equally striking:

One neuron, every input

Left: the selected neuron's activation for all (a, b). Right: the same grid in the 2D Fourier basis.

Loading…

How to read it
<div class="fig-read"><p>Hover the right panel to see which <a class="term" data-term="linear-term" href="../glossary.html#linear-term">linear term</a> or <a class="term" data-term="quadratic-term" href="../glossary.html#quadratic-term">quadratic term</a> a coefficient belongs to; click it to select that frequency.</p>

What to expect: A periodic pattern whose transform has only the 3×3 block of one frequency: constant, linear and quadratic terms.

Learn step by step → · Source: [1] Fig. 4 (centre); [3] “Neurons produce quadratic terms” · Computed by browser (GrokModel) (code)

The right panel is the neuron's 2D Fourier transform (chapter 10, with its average removed): almost all of it sits in the 3 × 3 block of one frequency — constant, linear terms and quadratic terms of that frequency. Neuron 0 of the paper's network is periodic with frequency 42 [1] (Sec. 4.1). Here are more neurons, three from each frequency cluster:

These neurons were picked from the final model's frequency clusters; played through training, each one turns from noise into stripes.

Neurons from every frequency cluster

Three neurons from each key-frequency cluster of the final model, shown at the selected step: play the run to watch the same neurons turn into waves.

Loading…

How to read it
<p class="fig-meta">
  <a href="../chapters/13-evidence-periodicity.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 12
  · Computed by <code>browser (GrokModel)</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Where do these waves come from? A neuron's input is what the attention heads copy from the tokens, passed through W_{in}. For one head, the map from the attended token a to all 512 neuron inputs is W_{neur} = W_{in} W_O W_V W_E (ARENA's name for it [3]). Written in the Fourier basis over a, each neuron's row is already concentrated on the waves of its own frequency — the periodicity is in place before the ReLU acts:

W_neur in the Fourier basis: what attention writes into each neuron

For the selected head, the map from the token a (attended to) to the 512 neurons' inputs, written in the Fourier basis over a.

Loading…

How to read it · the math
W_{neur}^{j} = W_{in}\,W_O^{j}\,W_V^{j}\,W_E \quad (512 \times 113)

Rows are neurons sorted by frequency cluster, columns the components Const, cos 1, sin 1, … over a. On the dashboard, pick the head in the control bar (chapters show head 0). This is ARENA's W_neur (paper: the OV circuit followed by W_in). It is computed from a's input vector at position 0, W_E + W_pos[0]; the position term only changes the Const column.

What to expect: Block structure: a neuron of cluster k receives only cos k and sin k of a, so each neuron's input is already a single wave before the ReLU.

Learn step by step → · Source: [3] “W_neur in Fourier basis” · Computed by browser (GrokModel) (code)

Logits are periodic

Finally the paper takes the 2D Fourier transform of the logits over (a, b) for every answer c and finds only twenty significant components — the products of sines and cosines at the five key frequencies, five 2 × 2 blocks (paper Fig. 4, right). We see the same in the dashboard's geometry and output panels, and test it quantitatively in chapter 15.

Nuance

Periodic ≠ explained. Periodicity shows which basis to look in, not what is computed with it. The next three chapters test the mechanism itself: neurons computing products (ch. 14), the logits reading off cos(w(a+b−c)) (ch. 15), and interventions that would break if the story were wrong (ch. 16).