Quick review

Every part of our model on one page, in the order the sentence "Lily wanted to play with her" passes through it. Each card has the shapes, and where it helps, the formula and a line of code.

Tokens in the example

T = 6

Numbers per token

d = 128

Blocks

L = 4

Heads per block

4

Numbers per head

d_k = 32

Feed-forward width

512

Vocabulary

4,096 tokens

Longest text

128 tokens

One forward pass gives every value on this page.

1. Inputs

   Your text, split into tokens. Each token is an id number from the vocabulary.

   The sentence "Lily wanted to play with her" is 6 token ids, each from 0 to 4,095.

   const ids = tokenizer.encode('Lily wanted to play with her');

   Read the chapter on InputsSee it on the map, Inputs

2. Embedding

   Each id picks one row from a table of learned numbers. That row is the token's starting list of numbers.

   The table has 4,096 rows of 128 numbers. The sentence becomes 6 rows of 128.

   $$\mathbf{x}_i = W_E[t_i] \in \mathbb{R}^{128}$$
   trace.tokenEmbeddings // 6 × 128

   Read the chapter on EmbeddingSee it on the map, Embedding

3. Position

   Each position also gets a list of numbers, added to the token's list. That is how the model knows the order. GPT-2 learns a table of positions instead.

   A table of 128 positions by 128 numbers. Row 0 goes to the first token, row 1 to the second.

   $$\mathbf{h}_i = W_E[t_i] + W_P[i-1]$$
   trace.embeddings // tokenEmbeddings + positionEmbeddings

   Read the chapter on PositionSee it on the map, Position

4. Add and norm

   Add, called a residual connection, keeps what came in and puts the new result on top, so earlier numbers are kept. Norm rescales each token's numbers to a steady range. GPT-2 moves the norm to the start of each step, before attention and before the feed-forward layer.

   Each token's 128 numbers are handled on their own.

   $$x + \mathrm{Sublayer}\big(\mathrm{LayerNorm}(x)\big)$$
   trace.layers[0].ln1.out // 6 × 128, normed before attention

   See it on the map, Add and norm

5. Masked attention

   Each token looks at itself and the tokens before it, and pulls in what it needs. Masked means it cannot look ahead.

   Our model has 4 heads with 32 numbers each. Each head makes a 6 by 6 grid of weights.

   $$\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V, \quad M_{ij} = -\infty \text{ when } j > i$$
   trace.layers[0].heads[0].weights // 6 × 6, each row sums to 1

   Read the chapter on Masked attentionSee it on the map, Masked attention

6. Feed forward

   Each token's numbers go through a small network on their own, without looking at other tokens.

   128 numbers grow to 512, pass through GELU, and shrink back to 128.

   $$\mathrm{GELU}(xW_1 + b_1)\,W_2 + b_2$$
   trace.layers[0].mlpOut // 6 × 128

   See it on the map, Feed forward

7. Repeated blocks

   Everything inside the frame is one block. The model repeats it, and each block refines the numbers further.

   Our model stacks 4 blocks. The paper used 6, and GPT-2 small uses 12.

   $$h^{(\ell+1)} = \mathrm{Block}_\ell\big(h^{(\ell)}\big)$$
   trace.layers.length // 4

   See it on the map, Repeated blocks

8. Final norm

   One more norm after the last block, before the scores are made. GPT-2 added it. The paper has none.

   The 6 rows of 128 numbers, rescaled.

   $$\mathbf{u} = \mathrm{LayerNorm}\big(h^{(L)}\big)$$
   trace.lnFinal.out // 6 × 128

   See it on the map, Final norm

9. Linear

   Turns each token's numbers into one score for every token in the vocabulary. It reuses the embedding table.

   4,096 scores per position. Only the last position is used to guess the next token.

   $$\mathbf{z} = \mathbf{u}\,W_E^\top$$
   trace.logits // 6 × 4,096

   Read the chapter on LinearSee it on the map, Linear

10. Softmax

    Turns the scores into chances between 0 and 1 that add up to 1.

    4,096 chances, one for each possible next token.

    $$p_t = \frac{e^{z_t}}{\sum_s e^{z_s}}$$
    probabilities(rowView(trace.logits, ids.length - 1), 1)

    Read the chapter on SoftmaxSee it on the map, Softmax

11. Output

    The model's guess, a chance for every possible next token. Pick one, add it to the text, and run again.

    One chance for each of the 4,096 tokens.

    $$P(t_{T+1} \mid t_1, \dots, t_T)$$
    sample(probs, createRng(1)) // the next token id

    Read the chapter on OutputSee it on the map, Output