A tiny GPT
Put the pieces together and you have a GPT: token embeddings plus position embeddings, a stack of blocks, a final normalisation, and a linear head that gives a score for every token in the vocabulary, at every position.
import torch
from torch import nn
tokens = nn.Embedding(20, 16)
positions = nn.Embedding(8, 16)
ids = torch.tensor([[3, 1, 4, 1]])
where = torch.arange(4)
x = tokens(ids) + positions(where)
print(x.shape) # torch.Size([1, 4, 16])Training shows it windows of text and asks, at every position, for the next token. Unlike the bigram model, it can use everything before, so after the c it knows at comes next even though t alone is ambiguous.
Your task
Write TinyGPT(vocab, d, context, blocks): nn.Embedding(vocab, d) named tokens, nn.Embedding(context, d) named positions, an nn.Sequential of blocks Block(d, 4 * d), an nn.LayerNorm(d) named norm and an nn.Linear(d, vocab) named head. forward(ids) takes (batch, T) and returns logits (batch, T, vocab). Write train_gpt(model, ids, context, steps=300, lr=3e-3): each step, take 16 random windows of context + 1 ids, train with Adam on cross-entropy of every next id, and return the list of losses as floats.
The explanation opens here once every check passes.