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

The task as a dataset

12,769 questions, 30 % of them for training, encoded as three tokens.

Sources for this chapter
  • [1] Sec. 3
  • [6] gen_train_test
  • [3] “Notation”

The question as a function

A function is a rule that gives exactly one output for each input. The rule the network must learn is

f(a, b) = (a + b) \bmod 113, \qquad a, b \in \{0, 1, \ldots, 112\}.

There are 113 choices for a and 113 for b, so there are 113 × 113 = 12,769 possible questions — few enough that we can check the network on every single one.

Examples, labels, and the training/test split

A dataset is a collection of examples: each example is an input together with its correct output, called its label. Here the example “(100, 50)” has label 37.

The paper's setup shows the network only a random 30 % of the questions while it learns [1]. Those form the training set: 3,830 questions. The remaining 8,939 questions form the test set: they are never used to adjust the network, only to check it.

Which inputs are used for training

All 113 × 113 inputs (a = row, b = column): yellow = training pair, purple = test pair.

Loading…

How to read it 
<p class="fig-meta">
  <a href="../chapters/02-task-as-data.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#paperrepo" title="progress-measures-paper code">[6]</a> gen_train_test
  · Computed by <code>task.paper_split</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The split is made by shuffling all 12,769 pairs with a fixed random seed and taking the first 30 % (0.3 × 12,769 = 3,830.7, rounded down) [6]. Hover the picture: each cell is one question (a = row, b = column); yellow cells are training questions.

Worked example

Is (100, 50) a training question? In the paper's split it is a test question — the network never sees its answer during training. (By contrast (5, 71) is a training question.) That is why (100, 50) is our running example throughout the site: when the trained network gets it right, it is not because it saw it.

Two ways to get training questions right

A network can answer its training questions correctly in two very different ways:

The test set tells the two apart: memorisation scores ~0 % on it (about 1 in 113 right by luck), generalisation scores ~100 %.

Questions as token sequences

The network reads its input as a short sequence of symbols called tokens, each written as an integer id. The question “100 + 50 =” becomes the three tokens

(\,100,\; 50,\; 113\,)

where ids 0–112 stand for the numbers and the extra id 113 stands for “=” [3]. The set of all tokens the network knows is its vocabulary — 114 tokens here. The position of a token is its place in the sequence: position 0 holds a, position 1 holds b, position 2 holds “=”.

Note that the “+” is never shown: every input is just “a b =”, and the network has to discover that the task is addition from the examples alone.

Nuance

The network outputs a score for each of 114 tokens (including “=”), but only the 113 number tokens are possible answers. Like the paper's analysis code, our analyses drop the “=” output [6]; training uses all 114, exactly as the paper's training code did.