Build a GPT forward pass

Learn it

Words to know

embedding

A list of numbers that stands for a token. The model learns these numbers during training.

positional encoding

Numbers added to a token's embedding to say where in the sequence the token stands.

causal mask

The rule that a token may attend only to itself and earlier tokens, never to later ones.

logit

A raw score a model gives to one choice, before the scores are turned into probabilities.

How it works

1. Two embedding tables of width d: one with a row per token, one with a row per position. Look up both and add them. Shape: [1, length, d].
2. Build the causal mask: a length by length table that is True above the diagonal. True means this position may not look there.
3. Pass the vectors and the mask through one transformer block: torch.nn.TransformerEncoderLayer(d, 2, dim_feedforward=2 * d, dropout=0, batch_first=True). The shape stays [1, length, d].
4. A linear layer from d to vocabulary_size gives the scores. Final shape: [1, length, vocabulary_size].
5. Check the mask: run a row of zeros, then the same row with its last token changed to 1. No earlier position can see the last token, so their scores must come out the same.

Worked example

Input
vocabulary_size = 5
length = 3
d = 4

Output
[[1, 3, 5], True]

One sequence, 3 positions, and 5 scores at each position. The position embeddings have no batch dimension; PyTorch adds the same rows to every sequence in the batch.

Why it matters

tiny_gpt.py has this forward pass with two changes: its block normalizes before each step, not after (norm_first=True), and one more normalization comes before the last layer. Its self-test runs the same check for looking ahead.

Your turn

Using PyTorch, build a one-block GPT. Return the shape of its output, and whether the scores at earlier positions stay the same when only the last token changes.

vocabulary_size

the number of tokens the model knows; at least 2

length

the number of positions in the sequence

d

the width of the embeddings; positive and even

Returns

The list [shape, unchanged]: the shape of the logits as a list, and True if a row of zeros and the same row with its last token set to 1 give equal logits at every position before the last.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
vocabulary_size = 5
length = 3
d = 4

Output
[[1, 3, 5], True]

Input
vocabulary_size = 6
length = 4
d = 4

Output
[[1, 4, 6], True]

Input
vocabulary_size = 7
length = 1
d = 8

Output
[[1, 1, 7], True]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   You need four modules: torch.nn.Embedding(vocabulary_size, d) for tokens, torch.nn.Embedding(length, d) for positions, a TransformerEncoderLayer, and torch.nn.Linear(d, vocabulary_size). Create them first, then pass the data through in that order.

2. Hint 2: The key step

   The positions are torch.arange(length), the numbers 0 to length - 1. The mask is torch.triu(torch.ones(length, length, dtype=torch.bool), diagonal=1). Put the forward pass in a small function, because you run it twice.

3. Hint 3: The outline

   Define forward(ids): add token(ids) and position(torch.arange(length)), run block(x, src_mask=mask), then head(...). Make first = torch.zeros((1, length), dtype=torch.long) and second = first.clone() with second[0, -1] = 1. For their outputs a and b, return [list(a.shape), torch.allclose(a[:, :-1], b[:, :-1], atol=1e-6)].

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. The four parts of the model. batch_first=True makes the block expect [batch, time, width]. dropout=0 turns off the random zeroing used in training. PyTorch calls this layer an encoder; the causal mask is what makes it a GPT block.
2. True above the diagonal: no looking at later positions. triu with diagonal=1 keeps the part above the diagonal. Row i is True in the columns after i.
3. Which token it is, plus where it stands: [1, length, d]. torch.arange(length) is the list of positions. Its embeddings have shape [length, d] and are added to every sequence in the batch.
4. Block, then one score per vocabulary token at each position. The block keeps the shape. The linear layer changes the last size from d to vocabulary_size.
5. Change only the last token: earlier scores must stay the same. a[:, :-1] is every position but the last. allclose compares the two with a small tolerance and returns a plain True or False.

The same numbered steps are the comments in the solution.

import torch


def solve(vocabulary_size, length, d):
    # 1. The four parts of the model
    token = torch.nn.Embedding(vocabulary_size, d)
    position = torch.nn.Embedding(length, d)
    block = torch.nn.TransformerEncoderLayer(
        d, 2, dim_feedforward=2 * d, dropout=0, batch_first=True
    )
    head = torch.nn.Linear(d, vocabulary_size)
    # 2. True above the diagonal: no looking at later positions
    mask = torch.triu(torch.ones(length, length, dtype=torch.bool), diagonal=1)

    def forward(ids):
        # 3. Which token it is, plus where it stands: [1, length, d]
        x = token(ids) + position(torch.arange(length))
        # 4. Block, then one score per vocabulary token at each position
        return head(block(x, src_mask=mask))

    # 5. Change only the last token: earlier scores must stay the same
    first = torch.zeros((1, length), dtype=torch.long)
    second = first.clone()
    second[0, -1] = 1
    a, b = forward(first), forward(second)
    return [list(a.shape), torch.allclose(a[:, :-1], b[:, :-1], atol=1e-6)]

Watch out

• Embedding indexes must be whole numbers: dtype=torch.long. A float tensor of IDs raises an error.
• In this mask True means blocked. With tril every position sees the tokens after it, and the check comes back False.

Check yourself