Attention + multi-head attention

Attention + multi-head attention, the mechanism

Where we are in the build

After Lesson 3, every position in your sequence is a (768)-dim vector that knows both what token it is and where it sits. Today's lesson is the one that made transformers transformers: every position looks at every other position and decides who matters, then mixes their information into its own representation.

We build it three times, in increasing realism:

1. Naked NumPy, no scaling, one head. Math you can read.
2. Scaled dot-product attention. Show why the $\sqrt{d_k}$ factor is not optional.
3. Multi-head PyTorch, n_head parallel attention computations sharing the same n_embd budget.

After this lesson, your scaffold has attention.py with a MultiHeadAttention module. We do not add the causal mask yet, that's Lesson 5. Without it the model can peek at future tokens, which is wrong for language modelling but useful for letting us see attention's behaviour in isolation.

Step 1, attention from scratch in NumPy

Forget projections, forget heads, forget scaling. Start with the simplest possible question: how does one position pull information from every other position?

Each position $i$ produces three things:

• A query $q_i$, "what am I looking for?"
• A key $k_i$, "what do I have to offer?"
• A value $v_i$, "what information do I carry?"

The score from position $i$ to position $j$ is the dot product $q_i \cdot k_j$, high if the query is aligned with that key, low otherwise. Softmax across all $j$ turns scores into a probability distribution. Position $i$'s output is the weighted sum of all values, weighted by that distribution.

That's the entire mechanism. Three matrices, one matmul for scores, a softmax, one matmul for the weighted sum. Every transformer in production today does this exact thing, with optimisations layered on top.

Step through it — press Run to execute the NumPy in your browser while the animation builds the Q·Kᵀ scores, the row-wise softmax heatmap, and the weighted-value output:

import numpy as npT, D = 4, 8x = np.random.randn(T, D)          # T tokens × D # project to queries, keys, valuesQ = x @ W_q                        # (T, D)K = x @ W_kV = x @ W_v # scores: q_i · k_j, scaledscores = Q @ K.T / np.sqrt(D)      # (T, T) # softmax row-wise → attention weightsweights = np.exp(scores) / np.exp(scores).sum(-1, keepdims=True) # weighted sum of valuesout = weights @ V                  # (T, D)

But there's a numerical bug hiding in those np.exp calls. Try it at full GPT-2 scale.

Step 2, why we scale by $\sqrt{d_k}$

Run the same code with $D = 64$ instead of $8$, the head dimension GPT-2 small actually uses:

T, D = 4, 64
W_q = np.random.randn(D, D)
W_k = np.random.randn(D, D)
 
x = np.random.randn(T, D)
Q = x @ W_q
K = x @ W_k
 
scores = Q @ K.T
print("score range:", scores.min(), scores.max())
 
weights = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
print("entropy of row 0:", -(weights[0] * np.log(weights[0] + 1e-10)).sum())

The scores blow up to numbers like $-150$ to $+200$. After softmax, each row puts ~99.99% of its mass on a single position. The "attention" is no longer a soft weighted sum, it's a hard argmax. Gradients through that softmax are essentially zero everywhere except the winning cell, which is the same vanishing-gradient problem we saw in Lesson 1.

Why does it happen? Each entry of $QK^T$ is a dot product of two vectors with $D$ components each. If those components are roughly $\mathcal{N}(0, 1)$, the dot product has variance $D$, its magnitude grows like $\sqrt{D}$. At $D = 64$ that's already in the tens, and softmax saturates fast.

The fix from Vaswani et al. (2017): divide by $\sqrt{d_k}$ before softmax, where $d_k$ is the per-head key dimension:

This rescales the dot products to roughly unit variance. The softmax stays in the soft regime. Gradients flow.

# One-line fix.
scores = Q @ K.T / np.sqrt(D)
 
weights = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
print("entropy of row 0:", -(weights[0] * np.log(weights[0] + 1e-10)).sum())
# entropy is back near log(T) = 1.39, the row is broadly distributed

That single division is why the original Transformer paper is titled "Attention Is All You Need" and not "Attention And A Numerical Hack You Need." The hack is built in.

Step 3, why multiple heads

A single attention layer with $D = 768$ has $3 \cdot 768^2 \approx 1.77$M parameters and produces a single attention pattern per layer. Empirically, that's both expensive and expressively limited, different "kinds" of relationships (subject-verb agreement, coreference, lexical similarity) want different attention patterns.

Multi-head attention runs $H$ smaller attention computations in parallel, each with its own $W_Q, W_K, W_V$ of dimension $C \to D$ where $D = C / H$. The outputs (each of dim $D$) are concatenated back to dim $C$ and passed through one more linear $W_O$. Same total compute and parameter count as a single head of dim $C$, but the model gets $H$ different attention patterns to combine.

For GPT-2 small:

The factor of $4 C^2$ comes from three projections ($W_Q, W_K, W_V$, each $C \times C$) plus the output projection $W_O$, also $C \times C$. The number does not depend on $H$: splitting one big $C$-dim head into $H$ smaller $D$-dim heads doesn't change the total weight budget, only how those weights are organised.

Code: attention.py

# nanogpt-from-scratch/attention.py  (Lesson 4, no causal mask yet)
 
import math
 
import torch
import torch.nn as nn
import torch.nn.functional as F
 
from config import GPTConfig
 
 
class MultiHeadAttention(nn.Module):
    """(B, T, n_embd) -> (B, T, n_embd)"""
 
    def __init__(self, config: GPTConfig):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_size = config.head_size           # C / H
 
        # One fused projection: (B, T, C) -> (B, T, 3C). Q, K, V at once.
        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
 
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        H = self.n_head
        D = self.head_size
 
        q, k, v = self.qkv(x).split(self.n_embd, dim=-1)
 
        # Reshape into heads: (B, T, C) -> (B, H, T, D) for parallel attention.
        q = q.view(B, T, H, D).transpose(1, 2)
        k = k.view(B, T, H, D).transpose(1, 2)
        v = v.view(B, T, H, D).transpose(1, 2)
 
        # Scaled dot-product: (B, H, T, T).
        att = (q @ k.transpose(-2, -1)) / math.sqrt(D)
        # NOTE: causal mask added in Lesson 5.
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
 
        y = att @ v                                  # (B, H, T, D)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
 
        return self.resid_dropout(self.proj(y))

Three things worth flagging:

• One fused QKV projection is faster than three separate Linear(C, C)s, same compute, better memory locality.
• bias=False on every Linear. GPT-2 turns biases off in attention; they don't help and they cost parameters. The GPT-2 paper and nanoGPT both do this.
• Two dropouts, one on the attention weights, one on the residual stream. We'll keep dropout=0.0 during pretraining (matches the GPT-2 paper) and only crank it up if we finetune later.

Click any query row to see where that token sends its attention across the keys, the live (T, T) heatmap this whole lesson is about:

Math by typing it

Single-head scaled dot-product attention:

with $Q, K \in \mathbb{R}^{T \times d_k}$ and $V \in \mathbb{R}^{T \times d_v}$. In GPT-2's self-attention, $d_k = d_v = D$ (the head dimension).

Multi-head attention runs $H$ copies in parallel and concatenates:

Why the variance argument matters: if $q, k \in \mathbb{R}^D$ have entries with mean 0 and variance 1, the dot product $q \cdot k = \sum_{j=1}^{D} q_j k_j$ has variance $D$. Standard deviation $\sqrt{D}$. Without scaling, softmax inputs grow with $\sqrt{D}$ and saturate. Dividing by $\sqrt{D}$ rescales to unit variance and keeps softmax soft.

Parameter count of one multi-head attention block (matching nanoGPT's no-bias setup):

For GPT-2 small ($C = 768$): $4 \cdot 768^2 = 2{,}359{,}296$ parameters per block. Across 12 blocks, ~28.3M total, roughly 23% of the model.

Run it on your laptop

import torch
 
from config import GPTConfig
from attention import MultiHeadAttention
 
config = GPTConfig()
attn = MultiHeadAttention(config)
 
# Fake input, pretend we already ran embeddings.
B, T, C = 2, 8, config.n_embd
x = torch.randn(B, T, C)
 
y = attn(x)
print("output shape:", y.shape)                 # torch.Size([2, 8, 768])
assert y.shape == x.shape, "attention should preserve (B, T, C)"
 
# Parameter check: should be exactly 4 * C^2.
n_params = sum(p.numel() for p in attn.parameters())
print(f"attention params: {n_params:,}")        # 2,359,296
assert n_params == 4 * config.n_embd ** 2
 
# Sanity: feeding the same input twice gives the same output (no randomness in eval).
attn.eval()
y1 = attn(x)
y2 = attn(x)
assert torch.allclose(y1, y2)
 
# Sanity: every output position depends on every input position.
# Perturb x at position 0 and see if y changes at position 5.
x2 = x.clone()
x2[:, 0, :] += 1.0
y_pert = attn(x2)
print("y changed at pos 5?", not torch.allclose(y[:, 5, :], y_pert[:, 5, :]))   # True

That last sanity check is important: without a causal mask, position 5's output does see position 0, and also position 6 and position 7. That's why we still need Lesson 5 to make this autoregressive.

What you can now do

• Compute scaled dot-product attention from raw Q, K, V tensors and explain every term in the softmax.
• Write the multi-head reshape pattern: (B, T, C) → (B, H, T, D) → attention → (B, T, C).
• Reason about the parameter budget, exactly $4 C^2$ per block, regardless of how many heads you slice it into.
• Diagnose softmax saturation when $d_k$ is large and explain why $\sqrt{d_k}$ scaling is the fix.

Quiz

Practice it yourself

Nail the nanoGPT specifics — the param budget, the fused split, and the √dₖ fix. Real Python, hidden tests, no setup.

What's next

You can build a multi-head attention block. But it has a fatal flaw for language modelling: at training time, position $t$ can attend to positions $t+1, t+2, \ldots$, the future. That means the model can "cheat" and copy the next token from its input. Lesson 5 fixes that with a single line: a triangular causal mask added to the attention scores before softmax. After Lesson 5, the model is finally autoregressive, and ready to be wrapped in a residual block in Lesson 6.