Grokking modular addition

Just the new ones. For earlier terms (feature, direction, dimension, superposition) see the step 1 glossary.

Grokking-specific

• Grokking: late, delayed generalisation. The phenomenon this project is about.
• Memorisation: the model perfectly fits training data without learning a rule. Looks like learning but doesn't transfer to new examples.
• Generalisation: the model performs well on examples it has never seen - evidence that it learnt the underlying rule, not the table of answers.
• Train / test split: a random partition of the data. The model only ever sees the train half during training; the test half is used to measure whether it actually learnt a generalising rule.
• Weight decay: a regularisation technique that pulls weights toward zero. Crucial for grokking - without weight decay, grokking does not happen. We'll explain why below.
• AdamW: the variant of the Adam optimiser that implements weight decay correctly. Just use it; you don't need to understand the details. (In code you'll see it spelt optimizer because PyTorch uses American spelling.)

Transformer anatomy

• Token: a single discrete input symbol. In our setup, the vocabulary is 114 tokens (0, 1, …, 112, =), and each input sequence has 3 tokens (a, b, =).
• Vocabulary (d_vocab): how many distinct tokens exist. For us, 114.
• Embedding (W_E): a learnable lookup table that maps each token to a vector of size d_model. The very first thing the model does to your input.
• Positional embedding (W_pos): a small learnable vector added per position in the sequence - lets the model tell apart "this is the first input" vs "this is the second input."
• Residual stream: the running vector at each position as it flows through the layers. The transformer's "main bus" - every component reads from it and writes back into it.
• Attention head: a component that lets the model mix information across positions. We have 4 heads in 1 layer.
• MLP (multilayer perceptron): a simple stack of linear layers with a nonlinearity in between (linear → ReLU → linear). Sits after attention and does per-position computation.
• Unembedding (W_U): a final linear layer that maps the d_model-sized vector back to d_vocab logits - one score per possible output token. The model picks the one with the highest score.
• Logits: the raw, unnormalised scores the model outputs. Softmax of the logits gives a probability distribution over tokens.
• Cross-entropy loss: the standard classification loss - penalises low predicted probability for the correct answer.

Mech interp / Fourier

• Fourier basis: a way of writing any function on {0, 1, ..., p-1} as a sum of sines and cosines at different frequencies. The mathematical structure the model ends up using.
• Frequency: an integer k indexing one of the trig functions cos(2πkx/p), sin(2πkx/p). With p=113, valid frequencies are 0 through 56.
• Progress measure: a quantity you can compute from the model's weights that smoothly increases even during the "flat" plateau. Reveals that the model is slowly building its circuit underneath, even when the test loss isn't dropping.

---

You did project 1, so you already know nn.Module, nn.Parameter, the training-loop skeleton, ReLU, loss, optimiser. The new concept here is the transformer.

A transformer is just a particular arrangement of standard neural-net pieces (linear layers, ReLU, softmax). Here's the entire architecture in one diagram, conceptually:

tokens (a, b, =)
        │
        ▼
   ┌──────────┐
   │ embedding│  W_E: lookup table that turns each token into a 128-dim vector
   └──────────┘  + positional embedding W_pos (per-position offset)
        │
        ▼
  [residual stream - a (3, 128) tensor: 3 positions × 128 features each]
        │
        ▼
   ┌──────────┐
   │ attention│  4 heads. Each head can move information between positions.
   └──────────┘  Writes its output back into the residual stream.
        │
        ▼
  [residual stream, updated]
        │
        ▼
   ┌──────────┐
   │   MLP    │  linear(128 → 512) → ReLU → linear(512 → 128)
   └──────────┘  Per-position computation.
        │
        ▼
  [residual stream, updated again]
        │
        ▼
   ┌──────────┐
   │unembedding│ W_U: linear(128 → 114). Produces a score for each possible answer.
   └──────────┘
        │
        ▼
   logits  →  softmax  →  predicted answer at the `=` position

Three concepts you need to internalise:

5.1 The residual stream

The transformer keeps a running (n_positions, d_model) tensor that flows through the layers. For us this is (3, 128) - three positions (one for a, one for b, one for =), each holding a 128-dim vector.

Every component (embedding, attention, MLP) reads from this tensor and adds its output back into it. Nothing is destructively overwritten. This is the "residual" part - outputs are added, not replaced.

Think of the residual stream as the model's working memory: a shared whiteboard that every component reads from and writes to.

5.2 Attention

An attention head's job is to move information between positions. In our setup, it's the only mechanism that lets the model combine information about a (at position 0) and b (at position 1) into a single computation at the = position (position 2). The MLP, by contrast, only operates within one position at a time.

Each head has four small matrices: W_Q (query), W_K (key), W_V (value), W_O (output). The mechanics:

1. For each position, the head computes a query (q = W_Q @ residual) and a key (k = W_K @ residual).
2. For each pair of positions, it computes an attention score: score[i, j] = q_i · k_j.
3. It softmaxes the scores along the source dimension to get attention weights (each position's "attention" sums to 1).
4. It computes values (v = W_V @ residual) and outputs a weighted sum of values, weighted by the attention weights.
5. It projects back via W_O and adds the result into the residual stream.

You don't have to memorise this. The mantra is: attention moves information between positions; everything else runs within a position.

5.3 Weight decay (and why it matters for grokking)

Weight decay is a regularisation technique: at every optimiser step, in addition to the gradient update, it pulls all weights slightly toward zero. Mathematically: w ← w - lr · gradient - lr · weight_decay · w.

Why does this matter for grokking? Roughly: there are two ways to solve the task.

• Memorisation: store a giant lookup table of (a, b) → answer in the weights. This works perfectly on the train set but uses a lot of weight "mass."
• Generalisation: learn a small, clean Fourier algorithm. Uses very little weight mass.

Without weight decay, the model is happy to sit at memorisation forever - both solutions get loss 0 on train, and there's no pressure to find the cleaner one. With weight decay, big-weight memorisation is penalised. The model has to slowly carve out the smaller, cleaner generalising solution. That carving is what we see as the long plateau, and grokking is the moment the cleaner solution becomes good enough to take over.

This is why the paper is called Progress Measures - even during the plateau, weight decay is slowly squeezing weight mass from the memorised solution into the generalising one. The progress is real but invisible in the loss curves until the crossover.

---