Build Your Own LLM, Lesson 6: Multi-Head Attention and the Transformer Block

📎 This article includes 4 downloadable practice files ↓

In Lesson 5 one attention head let each token look at earlier tokens. Real models use several heads at once, each free to learn a different kind of relationship (one tracks grammar, another names), then mix the results. Wrap that in a few more layers and you have the transformer block that GPT stacks dozens of times.

Multi-head attention

Instead of separate layers per head, compute queries, keys and values for all heads in one matrix multiply, then reshape:

q, k, v = self.qkv(x).split(C, dim=2)                       # (B, T, C) each
q, k, v = [t.view(B, T, n_heads, d_head).transpose(1, 2) for t in (q, k, v)]   # (B, heads, T, d_head)
att = (q @ k.transpose(-2, -1)) / math.sqrt(d_head)
att = att.masked_fill(mask[:, :, :T, :T] == 0, float("-inf"))   # causal: no peeking ahead
y = (F.softmax(att, dim=-1) @ v).transpose(1, 2).contiguous().view(B, T, C)

Each head works in a smaller space (d_model / n_heads), so four heads cost about the same as one big one.

The block

x = x + self.attn(self.ln1(x))      # attention, with a residual connection
x = x + self.mlp(self.ln2(x))       # feed-forward: Linear -> GELU -> Linear (4x wider inside)

• Residual connections (x + …) let information and gradients flow straight through deep stacks.
• LayerNorm before each part (“pre-norm”) keeps training stable.
• The MLP processes each token on its own after attention has mixed information between tokens.

Run it

Put tinygpt.py and the lesson script in one folder and run python lesson-06-transformer-block.py. Output shapes match the input (2, 8, 32), the causal check prints True, and one block has 12,608 parameters at this small size.

Practice

Run the script, then change n_heads to 8 and d_model to 64 and see how the parameter count changes.