Make your first transformer prediction

Make the vocabulary small enough to inspect

We will train on Tiny Shakespeare, a collection of Shakespeare excerpts. Here, one character is one token: letters, spaces, punctuation, and newlines each get an ID. That gives us 65 possible next tokens. We preserve capitalization and punctuation.

Lesson 1 used words as tokens. Character tokens let this small model assemble words one character at a time, with a small output table. Modern text LLMs commonly use subword tokens; changing the tokenizer changes the units, not the next-token task.

Our example input is ⏎ROMEO:⏎. The symbol ⏎ marks a newline, so this is eight characters. Each character becomes an integer that selects its embedding. The integer is an address, not a measure of meaning.

Follow the arrays through the model

PyTorch calls its arrays tensors. For now, follow one input. There is one row per character position; columns hold the numbers representing that position. We chose 64 numbers per representation; this width is separate from the 65-character vocabulary.

1. 8 token IDsIdentify the eight input characters.
2. 8 × 64 numbersLook up a 64-number token embedding and add a learned vector for its position.
3. 8 × 64 numbersTwo transformer blocks mix earlier information and process each position’s features.
4. 8 × 65 scoresThe output layer rates every possible next character at every position.

The model’s trainable numbers are called parameters, often also called weights. Each block keeps eight positions. A position may use itself and earlier positions, never later ones. Each block runs four attention heads side by side, each with its own learned comparisons, and combines their results. The second block can build on information combined by the first.

Every row predicts the character after its own position. To continue our input, generation uses only the last row: the 65 scores after the final newline. The other rows will become useful when we train on many next-character targets at once.

positions = torch.arange(token_ids.shape[1])
x = self.token_embedding(token_ids) + self.position_embedding(positions)
for block in self.blocks:
    x, weights = block(x)
logits = self.output(self.final_norm(x))