The whole Transformer

1. The whole machine

In order, from input to output: a token table turns each id into a vector. Then the position numbers of level 16 are added. Then come a stack of decoder blocks, one final LayerNorm, and an output layer that gives one score per token. Select a block to see its input, its output, and how many parameters it holds. The shapes are for one sequence of 8 tokens, the longest input this task has.

2. Under the microscope

The diagram shows the blocks. Now look inside them. Below, a tiny copy of this model (1 block, 1 head, d_model = 6, d_ff = 12) reads 8 9 = and writes the next word. Every size is different from every other size, so each number shows you which axis it is. Each step shows the tensor’s shape with named axes, its numbers, and the shape of the same tensor in the full model on this page.

Here is the block of the same tiny model as a picture, with the numbers from the microscope. The line underneath is the residual path: each sublayer reads a normalized copy of x and adds its result to x.

With several heads, attention scores have the shape (B, h, L_q, L_k): batch, heads, then the length of Q and the length of K. In a decoder-only model Q and K come from the same sequence.

3. One word at a time

Level 14 called it inference: running a trained model to get an answer. It is the whole forward pass, from token ids to logits, with every weight fixed. Training runs the same forward pass, then the backward pass, and then changes the weights.

The model never writes the whole answer at once. It writes one word, then runs again:

1. Start with the prompt: the digits and =.
2. Run the model over everything so far. The output has one row per token.
3. Only the last row matters. The output layer turns its 64 numbers into 42 scores, one per token. These raw scores are usually called logits.
4. Pick the highest score (this is called greedy decoding), and add that token to the end.
5. Repeat until the model picks <eos>.

Below is a real model, trained on this page’s task. Press Next step and watch each stage. The bars above the tokens show where the last row’s attention looked while it chose the next word.

Step 4 is the only new code. NumPy’s np.argmax gives the position of the largest value, not the value itself: np.argmax(np.array([0.2, 1.5, -0.3])) is 1.

Now the whole loop from the list above. The model returns one row of scores per token, so scores[-1] is the last row. ids.append(nxt) adds the picked id to the end, and break leaves the loop once the model picks <eos>. A fixed limit on new tokens stops a model that never picks <eos>.

One epoch is one pass over all 900 training numbers. The lab has the same model saved twice: after 2 epochs and after 60.

4. How it learned

While training, the model does not run one word at a time. It gets the whole sequence and makes every guess in one pass. The input x is the sequence without its last token, and the target y is the sequence without its first token:

427:    x   4  2  7  =     four     hundred  twenty  seven
        y   2  7  =  four  hundred  twenty   seven   <eos>

The causal mask stops each position from seeing the token it has to guess. Training on the correct previous tokens like this, instead of on the model’s own guesses, is called teacher forcing.

Not every guess is worth learning. Under the 4 the target is 2, and under the 2 it is 7: the model would be guessing the digits of the question, and those are random. Only the targets that are words or <eos> are the answer. Every target before them, = included, is part of the question. The loss counts the answer targets and skips the rest, by setting the skipped targets to PAD.

for x, y in batches:              # 1. a batch; skipped targets are PAD
    logits = model(x)             # 2. forward: (64, L, 42)
    # 3. loss, ignore_index=PAD
    loss = loss_fn(logits.reshape(-1, 42), y.reshape(-1))
    # 4. backward, then update
    opt.zero_grad(); loss.backward(); opt.step()