GPT From Scratch

A GPT is a character-level language model: token and position embeddings, a stack of blocks that each apply causal attention and a feed-forward network with residuals, a final normalization, and a linear head over the vocabulary. Nothing here is more exotic than the previous lessons.

The idea

Assembly, bottom to top:

• Embeddings. nn.Embedding(V, C) maps each token id to a vector; position embeddings (T, C) are added because attention cannot see order.
• Block. x = x + attn(ln1(x)), then x = x + mlp(ln2(x)). Normalizing before each sublayer is what makes deep stacks trainable.
• Residuals. Each x + is an identity path, so gradients reach the embeddings without a matmul.
• MLP. Linear(C, 4C) → GELU → Linear(4C, C), position by position.
• Head and loss. LayerNorm, then Linear(C, V) with no bias, scored by cross-entropy against the shifted input.

Worked example

Train a 2-layer, 2-head model with C = 32, T = 24, batch 16, and Adam at lr = 3e-3 on a 312-character corpus ("attention is all you need. " ×12) with 13 distinct characters.

• Parameters: 26,816
• Loss: 2.7146 at step 0, 0.0606 at step 100, 0.0506 at step 300, 0.0595 at step 599
• Step-0 loss is close to ln 13 = 2.565, the uniform guess over 13 characters

Sampling 60 characters at temperature 0.8 from "a" gives all you need. attention is all you need. attention is all you. The model memorized the phrase: spelling and long-range order from scratch, nothing about meaning.

In code

import torch, torch.nn as nn, torch.nn.functional as F

text = "attention is all you need. " * 12
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
V, C, T, B = len(chars), 32, 24, 16

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(C), nn.LayerNorm(C)
        self.attn = nn.MultiheadAttention(C, 2, batch_first=True, bias=False)
        self.mlp = nn.Sequential(nn.Linear(C, 4 * C), nn.GELU(), nn.Linear(4 * C, C))

    def forward(self, x):
        mask = torch.triu(torch.ones(x.size(1), x.size(1), dtype=torch.bool), 1)
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, C), nn.Embedding(T, C)
        self.blocks = nn.Sequential(Block(), Block())
        self.ln_f, self.head = nn.LayerNorm(C), nn.Linear(C, V, bias=False)

    def forward(self, idx, targets=None):
        x = self.blocks(self.tok(idx) + self.pos(torch.arange(idx.size(1))))
        logits = self.head(self.ln_f(x))
        if targets is None:
            return logits
        return F.cross_entropy(logits.view(-1, V), targets.reshape(-1))

torch.manual_seed(0)
model = GPT()
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
for step in range(600):
    i = torch.randint(0, len(data) - T - 1, (B,))
    xb = torch.stack([data[j:j + T] for j in i])
    yb = torch.stack([data[j + 1:j + T + 1] for j in i])
    loss = model(xb, yb)
    opt.zero_grad(); loss.backward(); opt.step()
    if step in (0, 100, 300):
        print(step, round(loss.item(), 4))       # 2.7146  0.0606  0.0506

print(sum(p.numel() for p in model.parameters()), round(loss.item(), 4))   # 26816 0.0595

model.eval()
idx = torch.tensor([[stoi["a"]]])
with torch.no_grad():
    for _ in range(60):
        idx = torch.cat([idx, torch.multinomial(F.softmax(model(idx[:, -T:])[:, -1] / 0.8, dim=-1), 1)], 1)
print("".join(chars[i] for i in idx[0].tolist()))

Check yourself

1. Why are position embeddings necessary when the model already sees T token embeddings in order?
2. What does each residual x + do for the gradient as it travels back to layer 0?
3. How many training signals does one causally-masked sequence of 24 tokens provide?

Key takeaways

• A GPT is embeddings plus pre-norm blocks of attention and MLP, wrapped in residuals.
• One causal forward pass over T tokens yields T next-token predictions, all trained.
• Generation runs the same network in a loop, sampling from the last position.