Architecture

Now we can assemble the modules from the previous chapter into the full transformer.

• forward computes the output by calling each module in order Embedding -> Attention -> MLP -> Linear (output layer).
• backward computes the gradient values by calling the same modules in the reverse order.
• predict runs the forward pass and computes the probability distribution over the vocabulary for the last token.
• train runs the full training loop calculating the loss first and then running the backward pass and updating the parameters.
• Attention and the MLP each use a residual connection. They add the output back to their input ($x + ...$, $r + ...$). It gives the gradient a direct path back through the network, which keeps training stable.
• Our training implementation runs one sequence at a time. Real LLMs process a batch of sequences at once to average the gradient over many examples so that the learning curve is smoother. It also parallelises better, but that doesn't matter for our small example.

The code is ready to be tested now. Let's train the model on our small corpus and run the generator predicting the next tokens for a given sequence.

Note that the corpus is very small in our example. The model will just memorize the exact order of the text.

The full code is available separately, so you can try it on a bigger corpus and at a larger scale.