Build Your Own LLM, Lesson 5: Self-Attention From Scratch

📎 This article includes 1 downloadable practice file ↓

🧠 Lesson 5 of 12 · Build Your Own LLM course

Our models so far predict the next token from one previous token. Real language needs more: in “The files that the manager sent were…”, the verb depends on files, five tokens back. Self-attention lets every token gather information from any earlier token — and learn which ones matter.

The intuition: a soft lookup

Think of each token asking a question and every earlier token offering an answer:

• Query (Q) — what am I looking for? (“I’m a verb; where is my subject?”)
• Key (K) — what do I contain? (“I’m a plural noun.”)
• Value (V) — what do I pass on if someone attends to me?

Each token compares its query with every key (a dot product), turns the scores into weights that add up to 1 (softmax), and takes a weighted average of the values. It is like an Excel lookup where, instead of one exact match, you get a blend of all rows weighted by how well they match.

Attention on four tokens with NumPy

import numpy as np
np.random.seed(0)

T, d = 4, 8                                   # 4 tokens, 8 numbers each
x = np.random.randn(T, d)                     # token embeddings (from Lesson 4)
Wq, Wk, Wv = (np.random.randn(d, d) / np.sqrt(d) for _ in range(3))

Q, K, V = x @ Wq, x @ Wk, x @ Wv              # each (4, 8)
scores = Q @ K.T / np.sqrt(d)                 # (4, 4): token i vs token j

mask = np.triu(np.ones((T, T)), k=1).astype(bool)
scores[mask] = -np.inf                        # no peeking at future tokens

weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True) # softmax per row
out = weights @ V                             # (4, 8): new token vectors
print(weights.round(2))

The printed weights form a lower-triangular matrix: row 1 can only attend to token 1, row 4 can attend to tokens 1–4, and every row sums to 1.

Three details that matter

The same thing in PyTorch

In a real model, Wq, Wk and Wv are learned by gradient descent, and we process a batch of sequences at once:

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

class Head(nn.Module):
    """One head of causal self-attention."""
    def __init__(self, d_model, d_head, block_size):
        super().__init__()
        self.q = nn.Linear(d_model, d_head, bias=False)
        self.k = nn.Linear(d_model, d_head, bias=False)
        self.v = nn.Linear(d_model, d_head, bias=False)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):                        # x: (batch, T, d_model)
        B, T, _ = x.shape
        q, k, v = self.q(x), self.k(x), self.v(x)
        att = q @ k.transpose(-2, -1) / k.shape[-1] ** 0.5      # (B, T, T)
        att = att.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        return att @ v                           # (B, T, d_head)

head = Head(d_model=32, d_head=16, block_size=64)
x = torch.randn(2, 10, 32)                       # 2 sequences, 10 tokens each
print(head(x).shape)                             # torch.Size([2, 10, 16])

Position: attention alone is order-blind

Attention compares every pair of tokens but has no idea which came first — “dog bites man” and “man bites dog” would look the same. Models fix this by adding a position embedding to each token embedding:

tok_emb = nn.Embedding(V, 32)
pos_emb = nn.Embedding(64, 32)                   # one vector per position
x = tok_emb(idx) + pos_emb(torch.arange(idx.shape[1]))

Many modern models use rotary position embeddings (RoPE) instead, but the learned version above is what GPT-2 used and is perfect for learning.

Why this changed everything

• Long context: a token can use information from anywhere in the window, not just its neighbour.
• Parallel training: every position is computed at once with matrix multiplications, which GPUs are extremely fast at. Older recurrent networks had to go one token at a time.
• Interpretable-ish: you can print attention weights and see what a token looked at.

Coming next

One head learns one kind of relationship. In Lesson 6 we run several heads side by side (multi-head attention), add a small feed-forward network, residual connections and layer normalisation — and that stack is a transformer block. Lesson 7 trains it on your own text.