Everything in sections 01 through 05 was preparation for this. We now read a minimal transformer implementation — every block, every line — and nothing will be unfamiliar. By the end, the black box is gone.
The code below is a minimal but complete transformer language model — roughly the architecture of GPT, scaled down to be readable. It can be trained on any text corpus and will learn to generate text in the same style. Every major language model in production today — GPT-4, Claude, Llama, Gemini — is a scaled and refined version of this structure.
We'll read it in blocks, mapping each section to the concepts you already understand.
Don't try to memorize this. Read it like documentation. Your goal is to be able to point to any line and say "I know what that's doing and why it's there." By the end of this section, you will.
# Hyperparameters — the knobs that define the model's size and behavior batch_size = 32 # how many sequences to train on simultaneously block_size = 256 # maximum context length (tokens per sequence) n_embd = 384 # embedding dimension — how wide each token's vector is n_head = 6 # number of attention heads per layer n_layer = 6 # number of transformer blocks (depth) dropout = 0.2 # fraction of neurons randomly disabled during training learning_rate = 3e-4 # step size for gradient descent max_iters = 5000 # total training steps vocab_size = 65 # number of unique tokens (characters in this case)
This is the model's specification sheet. Every number here corresponds to something we've discussed. n_embd is the width of the layer vectors from Section 03. n_head and n_layer are the attention and depth parameters from Sections 03 and 05. learning_rate is from Section 04. block_size is the context window — the quadratic attention cost bottleneck from Section 05.
In a production model like GPT-3, these numbers are: n_embd=12,288, n_head=96, n_layer=96, vocab_size=50,257. The architecture is identical. The scale is just dramatically larger.
# Build a character-level vocabulary from the training text with open('input.txt', 'r') as f: text = f.read() chars = sorted(list(set(text))) # all unique characters vocab_size = len(chars) # Two-way mapping: character ↔ integer stoi = { ch:i for i,ch in enumerate(chars) } itos = { i:ch for i,ch in enumerate(chars) } # Encode text to integers, decode integers back to text encode = lambda s: [stoi[c] for c in s] decode = lambda l: ''.join([itos[i] for i in l])
This is tokenization — converting text to numbers. This minimal version is character-level: each character becomes one integer. Production models use subword tokenization (BPE — Byte Pair Encoding), which groups common character sequences into single tokens. The principle is identical: every piece of text becomes a sequence of integers before the model ever sees it.
The vocabulary size (65 here, 50,257 in GPT-2) determines the size of the final output layer — the model must assign a probability to every token in the vocabulary at each step.
class Head(nn.Module): """One head of self-attention""" def __init__(self, head_size): # Three learned projections: query, key, value self.key = nn.Linear(n_embd, head_size, bias=False) self.query = nn.Linear(n_embd, head_size, bias=False) self.value = nn.Linear(n_embd, head_size, bias=False) # Mask: prevents attending to future tokens during training self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size))) def forward(self, x): B, T, C = x.shape # batch, time (sequence length), channels k = self.key(x) # (B, T, head_size) q = self.query(x) # (B, T, head_size) # Compute attention scores — every token queries every other token wei = q @ k.transpose(-2, -1) * C**-0.5 # Apply causal mask — can't look into the future wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) wei = F.softmax(wei, dim=-1) # Weighted aggregation of values v = self.value(x) out = wei @ v return out
This is the full attention head from Section 05, implemented. Every line maps directly to the concepts we covered.
The one new element is the causal mask — the triangular matrix of ones and zeros. During training, the model sees the correct next token (it's learning to predict it). The mask ensures it can't "cheat" by looking at future tokens — it can only attend to positions before the current one. This is called causal or autoregressive attention.
The * C**-0.5 is a scaling factor to keep the attention scores from getting too large before softmax — a numerical stability trick that the original paper introduced.
class MultiHeadAttention(nn.Module): def __init__(self, num_heads, head_size): self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)]) self.proj = nn.Linear(n_embd, n_embd) # recombine heads def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) return self.proj(out) class FeedForward(nn.Module): """Simple two-layer network applied to each token independently""" def __init__(self, n_embd): self.net = nn.Sequential( nn.Linear(n_embd, 4 * n_embd), # expand nn.ReLU(), nn.Linear(4 * n_embd, n_embd), # contract ) def forward(self, x): return self.net(x)
Multi-head attention is just multiple attention heads running in parallel, their outputs concatenated and projected back to the model dimension. Exactly as described in Section 05.
The FeedForward network is the other half of each transformer block. After attention allows tokens to share information with each other, the feed-forward network processes each token independently — expanding it to 4× the model dimension, applying ReLU, then contracting back. This is where much of the model's "thinking" happens per-token after context has been gathered.
class Block(nn.Module): """Transformer block: attention + feed-forward, with residual connections""" def __init__(self, n_embd, n_head): head_size = n_embd // n_head self.sa = MultiHeadAttention(n_head, head_size) self.ffwd = FeedForward(n_embd) self.ln1 = nn.LayerNorm(n_embd) self.ln2 = nn.LayerNorm(n_embd) def forward(self, x): x = x + self.sa(self.ln1(x)) # attention with residual x = x + self.ffwd(self.ln2(x)) # feed-forward with residual return x
The transformer block packages attention + feed-forward together with two important additions:
Layer Normalization (ln1, ln2): Normalizes the activations before each sub-layer. This keeps the training stable when stacking many layers — without it, values can grow or shrink uncontrollably as they pass through deep networks.
Residual connections (the x = x + ... pattern): The input is added back to the output of each sub-layer. This is a crucial trick for training deep networks — it gives gradients a "shortcut" path backwards through the network, preventing the vanishing gradient problem. You can think of each block as learning to add a small correction to the input rather than transforming it from scratch.
class GPTLanguageModel(nn.Module): def __init__(self): # Token embeddings: integer → vector self.token_embedding_table = nn.Embedding(vocab_size, n_embd) # Position embeddings: position → vector (so tokens know where they are) self.position_embedding_table = nn.Embedding(block_size, n_embd) # Stack of transformer blocks self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)]) # Final layer norm and projection to vocabulary self.ln_f = nn.LayerNorm(n_embd) self.lm_head = nn.Linear(n_embd, vocab_size) def forward(self, idx, targets=None): tok_emb = self.token_embedding_table(idx) # integers → vectors pos_emb = self.position_embedding_table(positions) # positions → vectors x = tok_emb + pos_emb # combine x = self.blocks(x) # through all transformer blocks x = self.ln_f(x) # final normalization logits = self.lm_head(x) # project to vocab size return logits # probabilities over vocabulary
The complete model in one class. Two things worth noting:
Position embeddings: Attention has no built-in sense of order — it looks at all tokens simultaneously. Position embeddings solve this by adding a learned vector to each token that encodes its position in the sequence. Without this, the model would treat "dog bites man" and "man bites dog" identically.
lm_head: The final linear layer projects from the model's internal dimension (n_embd=384) to the vocabulary size (65 here, 50,257 in GPT-2). Its output (the "logits") is the raw scores that get converted to probabilities via softmax. This is the layer that produces the "what's the next token?" answer.
# Training loop — the process from Section 04, implemented model = GPTLanguageModel() optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) for step in range(max_iters): # Get a random batch of training sequences xb, yb = get_batch('train') # Forward pass: compute predictions and loss logits, loss = model(xb, yb) # Backward pass: compute gradients optimizer.zero_grad() # clear gradients from previous step loss.backward() # backpropagation # Update weights optimizer.step() if step % 500 == 0: print(f"step {step}: loss {loss.item():.4f}")
This is Section 04, implemented verbatim. Forward pass, compute loss, backward pass, update weights. The optimizer here is AdamW — a more sophisticated version of gradient descent that maintains a running average of gradients to smooth out noisy updates. But the concept is identical to what we covered.
optimizer.zero_grad() is necessary because PyTorch accumulates gradients by default — you clear them at the start of each step to prevent them from compounding incorrectly across batches.
Configuration. Tokenization. Attention. Multi-head attention. Feed-forward. Transformer block. Full model. Training loop. That's the complete list. Everything else in a production model is optimization, safety, and scale — not new architectural concepts.