Chapter 21 of 0–21

The whole picture, again

The full objects, now fully explained — plus limitations, nuances, errata and related work.

Sources for this chapter

The objects, explained

Here are the three objects from chapter 0 again — now every piece has a name and a job.

1. What the network computes (chapters 5–6).

\underbrace{x^{(0)}_i = W_E t_i + W_{pos}[i]}_{\text{tokens} \to \text{vectors (embedding)}}\qquad \underbrace{A^j \approx \big(\sigma(C_j[a]{-}C_j[b]),\ 1{-}\sigma(\cdot),\ \approx 0\big)}_{\text{each head: a two-way choice between a and b}}

\underbrace{x^{(1)} = x^{(0)}_2 + \textstyle\sum_j W_O^jW_V^j \sum_i A^j_i x^{(0)}_i}_{\text{attention gathers a and b}}\qquad \underbrace{\mathrm{MLP}(a,b) = \ReLU(W_{in}x^{(1)} + b_{in})}_{\text{512 neurons; the ReLU multiplies}}

\underbrace{\Logits(a,b) \approx W_L\,\mathrm{MLP}(a,b)}_{\text{skip connection negligible; } W_L = W_UW_{out}}\qquad \text{answer} = \argmax_c\Logits(a,b)_c

And as a picture with the live numbers (click a stage for its formula):

One forward pass, stage by stage

The selected input flowing through the network at the selected step (embeddings, attention, residual stream, neurons, logits and the answer) with the real numbers.

Loading…

How to read it

<p class="fig-expect"><strong>What to expect:</strong> After training the correct answer's logit stands out and the prediction is right with probability close to 1.</p>
<p class="fig-meta">
  <a href="../chapters/05-the-model.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> App. A
  · Computed by <code>web/js/model/grokmodel.js</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

2. What the numbers mean (chapters 9–16).

\begin{aligned}&\underbrace{a \mapsto \cos(w_ka),\ \sin(w_ka)}_{\text{step 1: } W_E \text{ (sparse in Fourier, ch. 13)}}\\ \Rightarrow\;&\underbrace{\cos(w_k(a{+}b)),\ \sin(w_k(a{+}b))}_{\text{step 2: products via attention + ReLU (ch. 14)}}\\ \Rightarrow\;&\underbrace{\sum_{k\in\mathcal K}\alpha_k\cos(w_k(a{+}b{-}c))}_{\text{steps 3–4: } W_L \text{ reads off (ch. 15)}}\end{aligned}

For the paper's network 𝒦 = 14, 35, 41, 42, 52; the cosine formula explains 96.1 % of the logits; removing one key frequency multiplies the loss by 4,315× to 315,210×; swapping only the 10 key-frequency directions of a's embedding moves the answer 0.997 of the way to a′ + b.

The core of the algorithm on one clock:

3. How it got there (chapters 7, 8, 17–20).

\begin{aligned}&\theta \leftarrow \theta(1-\gamma\lambda) - \gamma\,\text{AdamStep}(\nabla_\theta\mathcal L):\\ &\underbrace{\text{memorisation}}_{\text{train loss}\downarrow} \;\to\; \underbrace{\text{circuit formation}}_{\text{restricted loss}\downarrow,\ \text{excluded loss}\uparrow} \;\to\; \underbrace{\text{cleanup}}_{\text{weights}\downarrow,\ \text{test loss}\downarrow\downarrow}\end{aligned}

On the paper's run the three phases end at steps 1,567, 9,301 and 14,001. The generalising circuit forms during the plateau, before the test loss drops; weight decay then removes the memorisation.

And the whole story in one view once more:

Grokking in one view

One moment of training seen four ways: the losses, the probability of the right answer on all 12,769 inputs, the numbers on one frequency's plane, and one neuron's activation. ▶ plays the run through grokking.

Loading…

How to read it

<div class="fig-read"><p>All four tiles show the same training step (the dotted line in the first tile). <strong>Tile 2</strong>: each pixel is one

question a + b (a = row, b = column); yellow means the network gives the right answer probability 1. Choose training pairs or test pairs to see the two groups separately. Tile 3: each dot is a number n, placed by projecting its embedding on the plane of one frequency's cosine and sine waves and coloured by n's position on that frequency's clock (k·n mod 113); a circle with the colours in order means the network stores n as an angle. The hands point at a, b and a + b. Tile 4: the selected neuron's activation on every input, recomputed in your browser from the nearest stored weights. While playing, tiles 2 and 3 blend the two stored frames around the current step (stored every 200 steps for the paper's run); paused, every tile shows a stored frame and its note says which.

What to expect: During memorisation tile 2 shows a speckle (training pairs right, test pairs wrong); the circle in tile 3 and the stripes in tile 4 form during the plateau; the whole grid turns yellow during cleanup.

Learn step by step → · Source: [1] Fig. 1–4 · Computed by export.py, geometry.py, web/js/model/grokmodel.js (code)

What we learned, in five sentences

A one-layer transformer trained on 30 % of all modular sums learns to represent numbers as angles at a few frequencies and adds them with the angle-addition identities ([1]; ch. 12).

Four lines of evidence — Fourier sparsity, single-frequency neurons, a rank-10 read-out, and interventions including counterfactual patching — support this description of the trained network (ch. 13–16).

Knowing the mechanism lets us measure hidden progress: the restricted loss falls and the excluded loss rises long before the test loss drops (ch. 17).

Throughout circuit formation and over most of memorisation the gradients of the training and test losses partly oppose each other (negative cosine, -0.44 to -0.08 in circuit formation), turning positive only during cleanup, after step 11,901; after grokking the network sits in a balance between the loss and weight decay (ch. 19).

In all 11 of our retrained seeds with weight decay, the Fourier signatures appear — sparse spectra, single-frequency neurons, cosine-shaped logits, and the patching test — each seed with its own key frequencies; the strictest preregistered ablation criterion (H4) held in only 6 of 11, because our rule for naming the key frequencies sometimes includes a weak one or leaves out a little used one. Without weight decay, none of 5 networks grokked within 40,000 steps (ch. 20).

Point 5 at a glance: the three phases of every run whose phases were detected, on one time axis, and the paper's run with its authors' own boundaries:

Every run's phases on one time axis

Memorisation, circuit formation and cleanup of every run whose phases were detected, with the step at which test accuracy passes 99 %; on the paper's row, ticks mark the paper's own boundaries.

Loading…

How to read it

<div class="fig-read"><p>Phases come from the objective rules of docs/operationalizations.md applied to each run's curves. The vertical cursor is the selected step; click a row to select that run.</p>

What to expect: The weight-decay seeds grok at different times, but always in the same order: memorisation, a long circuit formation, then cleanup.

Learn step by step → · Source: [1] Sec. 5.1 · Computed by phases.py (code)

Limitations and nuances

A toy model. One task, one layer, and a mechanism that took significant manual effort to find. The paper itself notes our progress metrics are specific to small networks on one particular algorithmic task (Nanda et al. 2023, Sec. 6, p. 9 · source)

Hindsight. The progress measures needed the final algorithm and do not predict when grokking happens (ch. 17).

Approximations. “Most” neurons are single-frequency (433 of 512 in the paper's run); the attention mechanism is described by the paper as a sketch; several components (e.g. always-firing neurons, the late attention change around step 43,000) remain partly unexplained (ch. 14, 18).

Reproduction. Our code reproduces most of the paper's numbers on its own network; a few intervention numbers could not be matched without knowing the exact procedure, and we list them openly (Reproduce, Errata).

Scope of our replication. We retrained only the paper's main setting; robustness to other data fractions, depths and moduli comes from the paper's appendix.

Related work

Grokking itself was discovered by Power et al. [2]; Liu et al. built an “effective theory” with a phase diagram of when it happens [16]; Thilak et al. linked loss spikes in adaptive optimisers to it [23].

Analytic solutions: Gromov derived the periodic weights of a two-layer MLP for modular arithmetic by hand [19].

Other algorithms: Zhong et al. found that modified architectures learn a different circuit (“pizza”) instead of this “clock” [17]; Chughtai et al. generalised the Fourier picture to arbitrary finite groups via representation theory [18].

Hidden progress: Barak et al. proposed hidden progress measures for learning parities [20].

Explanations elsewhere: Neel Nanda's original write-up [5], the ARENA course chapter [3], Google PAIR's interactive explorable with a smaller MLP [21], and the three videos [8] [9] [10].

Zhong et al.'s test can be run on this site's networks too: re-arranging the right answer's logit by the sum and the difference of the inputs gives their distance irrelevance, which separates a Clock from a Pizza.

Clock or pizza?

The logit of the right answer for every input, re-arranged by the inputs' sum a + b and difference b − a, its mean and spread at each difference, and the distance irrelevance q of Zhong et al.

Loading…

How to read it

<div class="fig-read"><p>Zhong et al. found two algorithms for modular addition. In the Clock (this site's story), the right answer's logit depends on the angles' sum only. In the Pizza, it also carries a factor |cos(w_k(a − b)/2)| that depends on the difference. Distance irrelevance q is the average spread of the right-answer logit among inputs with the same difference, divided by its spread over all inputs (their Def. 4.2): low q means the difference matters (their Pizza networks have q between 0 and 0.4).</p>

What to expect: For the trained networks of this site q stays well above the Pizza range (the live value is in the note): a Clock, as Zhong et al. report for transformers with full attention.

Learn step by step → · Source: [17] Sec. 4.1.2, Def. 4.2 · Computed by web/js/model/grokmodel.js (logitsGrid) (code)

Every term used on this site is in the glossary; every source in Sources.