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 8 of 0–21

Grokking

Memorise first, generalise thousands of steps later.

Sources for this chapter

The phenomenon

In 2022, Power and colleagues trained small transformers on tables of modular arithmetic and saw something odd [2]. The networks first fitted their training data perfectly while doing no better than chance on held-out data — and then, long afterwards, test performance jumped, in their words

“improving generalization performance from random chance level to perfect generalization”— Power et al. 2022, abstract, p. 1 · source

and

“this improvement in generalization can happen well past the point of overfitting”— Power et al. 2022, abstract, p. 1 · source

They called this grokking. Nanda et al. describe it as follows [1]:

“models abruptly transition to a generalizing solution after a large number of training steps, despite initially overfitting”— Nanda et al. 2023, Sec. 1, p. 1 · source

Here is the paper's own training run. Both curves are plotted on a logarithmic scale: each horizontal grid line is ten times the one below it, so we can see losses from ~10 down to ~10⁻⁷ at once.

Loss during training (log scale)

Training loss (blue) and test loss (red) at every 10th step; shaded bands are the three phases.

Loading…

How to read it · the math
\mathcal{L} = -\frac{1}{N}\sum_{(a,b)} \ln \softmax(\Logits(a,b))_{c^*}

Each curve is the average cross-entropy loss over its set of inputs. The vertical axis is a logarithmic scale, so each grid line is a factor of 10. Click anywhere to move the whole page to that step.

What to expect: Train loss falls during memorisation (until step 1,567 on the paper's run) while test loss rises, then, thousands of steps later, test loss collapses (grokking).

Learn step by step → · Source: [1] Fig. 2 · Computed by train.py / paper logs (code)

Worked example

Reading the plot (paper's run; numbers from its logs):

  • Steps 0 – 201: both losses start near ln 113 ≈ 4.73, the loss of uniform guessing; the paper's logs record 4.76 (train) and 4.77 (test). Training accuracy reaches 100 % by step 201.
  • 201 – 1,567: training loss keeps falling by orders of magnitude while test loss rises to about 26 — worse than guessing: the network confidently gives memorised-looking answers to questions it has never seen.
  • 1,567 – 9,301: a long plateau. Training loss is tiny and flat; test loss decreases only slowly.
  • Around 9,301: test loss collapses by seven orders of magnitude; test accuracy reaches 99 % at step 10,601.

Accuracy during training

Fraction of training (blue) and test (red) pairs answered correctly, at every checkpoint.

Loading…

How to read it · the math
\text{accuracy} = \frac{1}{N}\sum_{(a,b)} \mathbb{1}\big[\argmax_c \Logits(a,b)_c = c^*\big]
<p class="fig-expect"><strong>What to expect:</strong> On the paper's run training accuracy reaches 100 % by step <span class="num" data-num="paper.train_acc_full_step">201</span>; test accuracy stays low for thousands of steps but creeps up (<span class="num" data-num="paper.acc_test_501">5.6 %</span> at step 501, <span class="num" data-num="paper.acc_test_t1">7.6 %</span> at about step <span class="num" data-num="paper.t1">1,567</span>, <span class="num" data-num="paper.acc_test_7k">16.4 %</span> at step 7,001, <span class="num" data-num="paper.acc_test_t2">46.7 %</span> at step <span class="num" data-num="paper.t2">9,301</span>), then jumps to <span class="num" data-num="paper.acc_test_tgrok">99.5 %</span> by step <span class="num" data-num="paper.t_grok">10,601</span>.</p>
<p class="fig-meta">
  <a href="../chapters/08-grokking.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 2
  · Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The curves average over many questions. The next picture shows every question at once: each pixel is one of the 12,769 sums a + b (a = row, b = column), coloured by the probability the network gives to the correct answer (yellow = certain and right, dark purple = no chance). While it is on screen it plays the paper's run from the start to just after grokking; pause it or drag its slider to study one moment. Watch the training questions, a random speckle, turn yellow early while the test questions stay mostly dark through the plateau, until the whole square lights up after step 9,301.

Confidence on every input

Probability the model assigns to the correct answer, for all (a, b) at the selected step.

Loading…

How to read it 
<p class="fig-expect"><strong>What to expect:</strong> During memorisation the training pairs are bright and the test pairs dark (a speckle pattern = the random split); after grokking everything is bright.</p>
<p class="fig-meta">
  <a href="../chapters/18-watch-it-form.html">Learn step by step →</a>
  
  · Computed by <code>geometry.frame</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Why it looks sudden

On a linear axis the same data looks like an abrupt phase transition — nothing, then a cliff:

The same losses on a linear scale

Identical data with an ordinary (linear) vertical axis.

Loading…

How to read it 
<div class="fig-read"><p>On a linear axis the long plateau and the sudden drop look much more dramatic; the log plot shows that train loss keeps improving by orders of magnitude after it looks flat here.</p>

What to expect: A sharp cliff in test loss, much harder to see the slow changes before it.

Learn step by step → · Source: [8] (discussion of log vs linear axes) · Computed by train.py / paper logs (code)

The log plot already hints that something is changing during the plateau. Neel Nanda points out in his replication walkthrough that the choice of axis changes how sudden grokking appears [8]. The question this site answers, following the paper, is: what is the network doing during the plateau, and why does the test loss eventually collapse? The paper's answer in one sentence:

“grokking, rather than being a sudden shift, arises from the gradual amplification of structured mechanisms encoded in the weights”— Nanda et al. 2023, abstract, p. 1 · source

When does grokking happen?

Grokking in this setup depends on a few ingredients, all reported in the paper [1]:

In the paper's words:

“our networks do not grok on the modular arithmetic task without weight decay or some other form of regularization”— Nanda et al. 2023, Sec. 5.3, p. 9 · source

We test the weight-decay claim ourselves in chapter 20, with five runs without weight decay.

Nuance

Not “after zero loss”. Grokking is sometimes summarised as “the model generalises after reaching zero training loss”. As chapter 4 showed, cross-entropy is never zero; the training loss keeps decreasing throughout the plateau, and the paper explicitly warns against this summary (footnote 4) [1].