Multi-head attention

Two heads, two jobs

Here is a hand-built example with six tokens: "the rain jacket is light packable". Each token's vector has 9 numbers. The first 6 mark its position, with a 1 in its own slot and 0 everywhere else. This is called one-hot. The last 3 mark its category: product word, feature word or filler.

The projections are set by hand here, not learned. Head 1's make every token look for the previous position. Head 2's make every token look for tokens of the same category. Both heads use a causal mask, so a token can only look at itself and earlier tokens. Same input, two different patterns:

import numpy as np
import matplotlib.pyplot as plt

tokens = ["the", "rain", "jacket", "is", "light", "packable"]
category = [2, 0, 0, 2, 1, 1]                     # 0 = product word, 1 = feature, 2 = filler
T = len(tokens)
X = np.hstack([np.eye(T), np.eye(3)[category]])   # 6 position features + 3 category features

def attention(Q, K):
    s = Q @ K.T
    s = np.where(np.tril(np.ones((T, T))) == 1, s, -np.inf)   # causal mask
    w = np.exp(s - s.max(axis=1, keepdims=True))
    return w / w.sum(axis=1, keepdims=True)

shift = np.eye(T, k=-1) * 6                        # query for position i asks for position i - 1
W_q1 = np.vstack([shift, np.zeros((3, T))])
W_k1 = np.vstack([np.eye(T), np.zeros((3, T))])
W_q2 = np.vstack([np.zeros((T, 3)), np.eye(3) * 4])   # query and key: just the category
W_k2 = W_q2

heads = {"head 1: previous token": attention(X @ W_q1, X @ W_k1),
         "head 2: same category": attention(X @ W_q2, X @ W_k2)}
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
for ax, (name, W) in zip(axes, heads.items()):
    ax.imshow(W, cmap="viridis")
    ax.set_xticks(range(T)); ax.set_xticklabels(tokens, rotation=45)
    ax.set_yticks(range(T)); ax.set_yticklabels(tokens)
    ax.set_title(name)
    print(f"{name}: 'packable' attends most to {tokens[int(W[-1].argmax())]!r}")

Each picture is a grid of weights. Each row is a token, and the bright squares show where it looks. Head 1 makes a bright stripe just left of the diagonal: every token looks at the one before it. Head 2 lights up tokens that share a category.

Both lines print "light", for different reasons. In head 1, "packable" puts about 99% of its weight on "light", the word before it. In head 2, it splits its weight half and half between the two feature words, "light" and itself. (When two weights tie, argmax picks the first, so it names "light".)

category = [2, 0, 0, 2, 1, 1]     # the, rain, jacket, is, light, packable
i = 2                             # "jacket"
same = [j for j in range(i + 1) if category[j] == category[i]]
print(same)

Heads that learn their own jobs

In a trained model, nobody designs these patterns. Heads specialize because it helps the model predict the next token. Researchers have found heads that track the previous token, copy repeated patterns, match brackets, or link pronouns such as "she" to names.

The friends with highlighters were told their jobs. Real heads aren't. Their jobs come out of training, and they are often messier than this example, with several heads overlapping.

The real module

In real code, all the heads run together in one go. The trick is reshaping: changing how the same numbers are grouped, without changing the numbers.

The code names the input's shape with three letters. B is the batch, how many sequences are processed at once. T is the number of tokens in each sequence. C is the width, $d_{\text{model}}$ (C stands for channels). The steps are:

1. One linear layer makes the queries, keys and values for every head at once. Each comes out with shape (B, T, C).
2. view cuts the last dimension into (heads, head size). Then transpose(1, 2) moves the heads next to the batch, giving (B, h, T, d_head). Now each head looks like a separate item in the batch.
3. Attention runs once, on all heads together. Each head still gets its own scores and its own softmax.
4. transpose and reshape undo the split. That's the concatenation.
5. self.out is $W_O$.

import math
import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.d_head = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)   # Q, K and V for all heads in one matmul
        self.out = nn.Linear(d_model, d_model, bias=False)       # W_O

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)                   # each (B, T, C)
        q, k, v = (t.view(B, T, self.h, self.d_head).transpose(1, 2) for t in (q, k, v))  # (B, h, T, d_head)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)                            # (B, h, T, T)
        scores = scores.masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        y = torch.softmax(scores, dim=-1) @ v                                                # (B, h, T, d_head)
        y = y.transpose(1, 2).reshape(B, T, C)                   # concatenate the heads: (B, T, C)
        return self.out(y)

torch.manual_seed(0)
mha = MultiHeadAttention(d_model=32, n_heads=4)
x = torch.randn(2, 6, 32)
print("output:", tuple(mha(x).shape))
print("parameters:", sum(p.numel() for p in mha.parameters()), "= 4 × 32 × 32 (Q, K, V and output), whatever the number of heads")

The output shape, (2, 6, 32), matches the input: 2 sequences, 6 tokens, width 32.

Here torch is the course's in-browser teaching version of PyTorch. The code is the same as real PyTorch, but it runs only on your computer's main processor and much more slowly. Seeded random numbers also differ from real PyTorch.

Counting the parameters

The module printed 4,096 parameters. Splitting into heads didn't add any. Here's why.

Each head's query projection turns $d_{\text{model}}$ numbers into $d_{\text{head}}$ numbers. So it holds $d_{\text{model}} \times d_{\text{head}}$ weights. Its key and value projections are the same size. Then $W_O$ is a full square: $d_{\text{model}} \times d_{\text{model}}$ weights. (This module has no bias terms.)

In plain words: $h$ heads times the head size is the full width again. So the heads' projections add up to three full-width squares, plus one more square for $W_O$. The number of heads cancels out.

Work it through for the module above: width 32, 4 heads, head size 8.

• One head's three projections: 3 × 32 × 8 = 768.
• Four heads: 4 × 768 = 3,072.
• Add $W_O$: 32 × 32 = 1,024.
• Total: 3,072 + 1,024 = 4,096, which is 4 × 32 × 32.

The same count in Python, for several head counts:

d_model = 32
for h in [1, 2, 4, 8]:
    d_head = d_model // h
    per_head = 3 * d_model * d_head               # Q, K and V projections for one head
    total = h * per_head + d_model * d_model      # all heads, plus W_O
    print(h, "heads of size", d_head, "->", total, "parameters")

Every line prints 4096. What changes with more heads is the number of softmaxes. Four heads of size 8 give four separate attention patterns instead of one.

import numpy as np

B, T, C, h = 1, 3, 4, 2                     # 1 sequence, 3 tokens, width 4, 2 heads
x = np.arange(12.0).reshape(B, T, C)        # token 0 = [0, 1, 2, 3], token 1 = [4, 5, 6, 7], ...
heads = x.reshape(B, T, h, C // h).transpose(0, 2, 1, 3)    # split: (B, h, T, d_head)
print("split shape:", heads.shape)
print("head 0 sees:", heads[0, 0].tolist())
print("head 1 sees:", heads[0, 1].tolist())
merged = heads.transpose(0, 2, 1, 3).reshape(B, T, C)       # merge: (B, T, C)
print("merged equals x:", np.array_equal(merged, x))

Sizes in real models

GPT-2 small uses 12 heads of 64 numbers each, for a width of 768. Large language models use dozens of heads in each layer, across dozens of layers.

Many recent models go one step further. Several query heads share the same keys and values. This is called grouped-query attention, and it cuts the memory needed while the model generates text. That memory is the subject of the KV cache lesson.