Lesson 12 of 20Assembling the Transformer

Multi-Head Attention

The transformer block assembled last lesson used single-head attention: one set of Q/K/V projections, capturing one kind of relationship. But a sentence has many kinds of relationships at once, which word is the subject, which pronoun refers to which noun, which word disambiguates another. A single head has to compromise across all of them.

The fix: run several attention computations in parallel, multiple heads, each with its own learned Q/K/V projections, so each head is free to specialize. Their outputs get concatenated and mixed back together with one more projection.

Two heads, same query, different weight vectors — click a token

Head A (dim 0)

The

cat

sat

the

mat

because

was

tired

Head B (dim 1)

The

cat

sat

the

mat

because

was

tired

Concatenate both heads' outputs, mix with an output projection W_O: [1.02, 0.38]

Head A and Head B here use different (fixed, toy) projections, and you can see them genuinely disagree about what a token should attend to, that's the entire point. A real model might have 8, 12, or dozens of heads, each nudged by training toward noticing something different, and it's the combination of all of them, mixed by the output projection, that gives a transformer block its expressive power.

Python

import math

tokens = ["The", "cat", "sat", "on", "the", "mat", "because", "it", "was", "tired"]
embeddings = [
    [0.1, 0.1], [1.5, 0.2], [0.3, 1.2], [-0.2, 0.4], [0.1, 0.1],
    [-0.3, -1.0], [0.6, -0.6], [1.4, 0.3], [0.2, -0.3], [1.0, -0.8],
]

# each head projects the 2D embedding down to a 1D score space with its
# own weight vector, so different heads key off different dimensions
heads = [
    {"q": (1, 0), "k": (1, 0), "v": (1, 0)},  # head A: dimension 0
    {"q": (0, 1), "k": (0, 1), "v": (0, 1)},  # head B: dimension 1
]
W_O = [[1, 0.3], [0.3, 1]]

def project(v, w):
    return v[0] * w[0] + v[1] * w[1]

def softmax(scores):
    biggest = max(scores)
    exps = [math.exp(s - biggest) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]

def head_output(head, query_index):
    q = project(embeddings[query_index], head["q"])
    scores = [project(e, head["k"]) * q for e in embeddings]
    weights = softmax(scores)
    v_vals = [project(e, head["v"]) for e in embeddings]
    return weights, sum(w * v for w, v in zip(weights, v_vals))

query = tokens.index("it")
head_outputs = []
for name, head in zip(["A", "B"], heads):
    weights, out = head_output(head, query)
    head_outputs.append(out)
    print(f"head {name} top attention: {tokens[weights.index(max(weights))]} ({max(weights):.3f})")

combined = [
    W_O[0][0] * head_outputs[0] + W_O[0][1] * head_outputs[1],
    W_O[1][0] * head_outputs[0] + W_O[1][1] * head_outputs[1],
]
print("combined output:", [round(x, 3) for x in combined])

import math

tokens = ["The", "cat", "sat", "on", "the", "mat", "because", "it", "was", "tired"]
embeddings = [
    [0.1, 0.1], [1.5, 0.2], [0.3, 1.2], [-0.2, 0.4], [0.1, 0.1],
    [-0.3, -1.0], [0.6, -0.6], [1.4, 0.3], [0.2, -0.3], [1.0, -0.8],
]

# each head projects the 2D embedding down to a 1D score space with its
# own weight vector, so different heads key off different dimensions
heads = [
    {"q": (1, 0), "k": (1, 0), "v": (1, 0)},  # head A: dimension 0
    {"q": (0, 1), "k": (0, 1), "v": (0, 1)},  # head B: dimension 1
]
W_O = [[1, 0.3], [0.3, 1]]

def project(v, w):
    return v[0] * w[0] + v[1] * w[1]

def softmax(scores):
    biggest = max(scores)
    exps = [math.exp(s - biggest) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]

def head_output(head, query_index):
    q = project(embeddings[query_index], head["q"])
    scores = [project(e, head["k"]) * q for e in embeddings]
    weights = softmax(scores)
    v_vals = [project(e, head["v"]) for e in embeddings]
    return weights, sum(w * v for w, v in zip(weights, v_vals))

query = tokens.index("it")
head_outputs = []
for name, head in zip(["A", "B"], heads):
    weights, out = head_output(head, query)
    head_outputs.append(out)
    print(f"head {name} top attention: {tokens[weights.index(max(weights))]} ({max(weights):.3f})")

combined = [
    W_O[0][0] * head_outputs[0] + W_O[0][1] * head_outputs[1],
    W_O[1][0] * head_outputs[0] + W_O[1][1] * head_outputs[1],
]
print("combined output:", [round(x, 3) for x in combined])

Why do transformers use multiple attention heads instead of one larger one?

Three lessons of pieces, now put them all in one place. Pick any token below and step through (or hit Play) to watch its real, computed vector move through every stage of one complete transformer block.

The transformer block, assembled, click a stage or hit play

→

→× N blocks

Pick a token above, then click any stage (or hit Play) to see its real computed vector and what that stage actually does.

The block output is still just a vector per token, a richer, context-aware one, but not yet a prediction. Turning it into an actual next-token probability, and training the whole stack to get better at that, is next.

Go deeper

The Illustrated GPT-2 (Visualizing Transformer Language Models)Jay Alammar

Visualizes exactly the block you just assembled, stacked self-attention and feed-forward layers, inside a real, complete GPT-2 model.

Let's build GPT: from scratch, in code, spelled out.Andrej Karpathy

Builds and trains a real GPT-style Transformer in code from the ground up, the natural next step after assembling one block by hand.

← Residuals, Normalization & the Feed-Forward Layer Training: From Logits to Loss →