Chapter 5 of 0–21
The model, step by step
A one-layer transformer computing one answer, followed number by number.
What a neural network is
A neural network is a function assembled from simple steps — mostly matrix–vector products and a few simple nonlinear functions — whose numbers are not written by hand but tuned from data. Each tunable number is a parameter (weight) (also called a weight). Some parameters are a bias: a number added regardless of the input.
The network studied here is a transformer, the architecture behind today's language models [12], shrunk to a single layer (one attention step followed by one MLP step) [1]. It has 226,816 parameters. Running it on an input from start to finish is a forward pass. Below we follow one forward pass on our running example (a, b) = (100, 50); the widget computes every intermediate number with the paper's trained weights.
The same pass as a picture: every stage with its numbers for the selected question, and small pulses running from the tokens to the answer. While you read the four steps below, the picture stays in view and lights up the stage being explained, with its formula; click a stage to see another. When the picture comes into view it also plays the network's training once (chapters 7–8): 100 + 50 is a question the network never trained on, answered wrongly at first and correctly, with near certainty, at the end. ▶ replays it.
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>
Step 1 — embeddings: tokens become vectors
Each token is turned into a vector of dimension d = 128 by its embedding: column t of the matrix W_E (via the one-hot product of chapter 3). A positional embedding — row i of W_{pos} — is added so the network knows where each token sits:
For (100, 50, =): x^{(0)}_0 = column 100 of W_E + position-0 vector, x^{(0)}_1 = column 50 + position-1 vector, x^{(0)}_2 = column 113 + position-2 vector.
These three vectors start the residual stream: a running vector at each position that every later step reads from and adds its result to [11].
Step 2 — attention: the “=” position gathers information about a and b
Attention lets each position read from positions at or before it. The answer is read at the “=” position (position 2), and only its result matters, so we follow position 2. The model has H = 4 attention heads, each with its own four matrices. Head j:
- computes a query for position 2, q = W_Q^j x^{(0)}_2 (32 numbers), and a key for every position i, k_i = W_K^j x^{(0)}_i;
- computes an attention score s_i = q\cdot k_i/\sqrt{32} for each position — the dot product measures how well the key “matches” the query; dividing by \sqrt{32} keeps the numbers in a reasonable range [12];
- turns the three scores into the attention pattern A^j = \softmax(s_0, s_1, s_2) — three weights that sum to 1;
- computes a value v_i = W_V^j x^{(0)}_i for each position, mixes them with the pattern, and writes the result back to dimension 128 with W_O^j:
The causal mask forbids positions 0 and 1 from reading later positions (it sets those scores to a huge negative number, so their softmax weight is 0); position 2 may read all three. Two matrix products summarise a head [11]: the QK circuit W_Q^{j\top}W_K^j decides where to look, and the OV circuit W_O^jW_V^j decides what gets written. The outputs of the four heads are added to the residual stream:
Attention in miniature. Suppose head 0's scores at the “=” position are s = (1.0, 0.6, -5) for (a, b, =). Softmax gives weights ≈ (0.598, 0.401, 0.001): the head reads 60 % from a, 40 % from b and almost nothing from “=”. If a's value vector after W_O is u_a and b's is u_b, the head writes ≈ 0.598\,u_a + 0.401\,u_b. For the trained network the weights to “=” are indeed tiny (about 0.1–0.4 %) [1].
Step 3 — the MLP: 512 neurons
The MLP multiplies x^{(1)} by W_{in} (512 × 128), adds the bias b_{in}, and applies the ReLU, \ReLU(z) = \max(0, z), to each of the 512 entries. Each entry is a neuron: its value before the ReLU is its pre-activation, after the ReLU its activation.
ReLU on four pre-activations (2.3, −0.7, 0.0, 5.1) gives activations (2.3, 0, 0, 5.1). It is the only nonlinearity in the MLP, and — chapter 14 shows — the network uses it to multiply two waves together.
Step 4 — unembedding: one logit per answer
Finally the unembedding matrix W_U turns the 128 numbers of x^{(2)} into 113 logits (the 114th, for “=”, is ignored), and the answer is the argmax:
All the parameters
| Matrix | Shape (as stored) | Count | Role |
|---|---|---|---|
| W_E | 128 × 114 | 14,592 | token embeddings |
| W_{pos} | 3 × 128 | 384 | positional embeddings |
| W_Q, W_K, W_V | 3 × (4 × 32 × 128) | 49,152 | queries, keys, values of 4 heads |
| W_O | 128 × 128 | 16,384 | head outputs |
| W_{in}, b_{in} | 512 × 128, 512 | 66,048 | neuron inputs |
| W_{out}, b_{out} | 128 × 512, 128 | 65,664 | neuron outputs |
| W_U | 128 × 114 | 14,592 | unembedding |
| Total | 226,816 |
There is no layer normalisation and the embedding and unembedding are separate matrices, as in the paper [1]; the only biases are those of the MLP [6].