Multi-head attention: Visual walkthrough — AI-ML deep dive

We will re-derive the parent's central result: $MultiHead (Q, K, V) = Concat (head_{1}, \dots, head_{h}) W^{O}$ but slowly, in pictures.

Prerequisites we lean on (open in another tab if unfamiliar): Query-Key-Value Intuition, Self-Attention, Linear Projections in Neural Networks.

Step 1 — A sentence is a stack of arrows

WHAT. Before any attention happens, each word is already a list of numbers — a vector. Picture a vector as an arrow in space, or more practically as a short horizontal row of numbers. Put one row per word and you get a matrix: a grid of numbers, rows = words, columns = features.

WHY start here. Attention never sees letters. It sees this grid. Everything downstream is "reshaping the grid", so we must see the grid first.

PICTURE. Below, the sentence "the cat sat" becomes 3 rows. Each row has $d_{model}$ numbers (here we draw only 6 columns to fit the page; the real value might be 512).

In self-attention the same grid $X$ plays all three roles $Q$ , $K$ , $V$ — we will explain those next.

Step 2 — Three questions: Query, Key, Value

WHAT. From the one grid $X$ we make three new grids by multiplying it with three learned number-grids (this multiplication is a linear projection — see Linear Projections in Neural Networks): $Q = X W^{Q}, K = X W^{K}, V = X W^{V}$

WHY three? Each word must do three different jobs:

ask "what am I looking for?" → its Query row,
advertise "what do I offer?" → its Key row,
carry "here is my actual content" → its Value row.

One grid can't hold three different jobs cleanly, so we learn three separate reshaping matrices $W^{Q}, W^{K}, W^{V}$ .

PICTURE. Term by term, the multiplication turns each word-row into a query-row, a key-row and a value-row.

Why learned and not fixed? Because the network discovers, during training, which mix of features makes a useful question. See Query-Key-Value Intuition.

Step 3 — One head, one dot product = "how much do I care?"

WHAT. To measure how much word $i$ should attend to word $j$ , we compare word $i$ 's query row with word $j$ 's key row using a dot product: multiply matching entries and add them up.

$score_{ij} = q_{i} \cdot k_{j} = \sum_{c} q_{i, c} k_{j, c}$

WHY the dot product and not, say, subtraction? The dot product is large when two arrows point the same way and small (or negative) when they point apart. It is the cheapest possible "are these aligned?" meter — exactly the question "does this key answer this query?".

PICTURE. Two rows lined up, multiplied entry-by-entry, summed to one number. A big number = strong alignment = a bright beam.

Doing this for every pair fills a whole score grid $Q K^{T}$ of shape $n \times n$ .

Step 4 — Scale, then softmax = turn scores into a fair split of attention

WHAT. We divide every score by $d_{k}$ , then push each row through softmax.

$A = softmax (\frac{Q K ^{T}}{d _{k}})$

Softmax takes a row of numbers and turns it into positive weights that add up to 1 — a probability spread.

WHY divide by $d_{k}$ ? When rows are long ( $d_{k}$ big), dot products grow large just from having more terms to add. Huge scores make softmax spike to almost-all-on-one-word, and gradients die. Dividing by $d_{k}$ (the standard deviation of a sum of $d_{k}$ random terms) shrinks them back to a healthy range. Here $d_{k}$ = the width of one head's query row.

WHY softmax and not just "pick the max"? We want a soft blend — word 3 can be 45% about itself and 12% about word 1. Hard-picking would throw away the rest. Softmax is smooth, so the network can learn through it.

PICTURE. One row of raw scores → same row scaled → same row as bars summing to 1.

Edge case — a word attending only to itself. If one score dwarfs the others, softmax sends that weight toward 1 and the rest toward 0. Nothing breaks; the word simply copies its own value. Edge case — all scores equal. Softmax returns a flat $1/ n$ each: the word averages everyone equally. Both extremes are valid, and softmax slides smoothly between them.

Step 5 — Weighted sum of Values = the head's answer

WHAT. Multiply the attention weights by the value grid: $head = A V$ Row $i$ of the output is a blend of all value rows, weighted by how much word $i$ attended to each.

WHY use $V$ and not $K$ again? Keys were only for matching. Once we know how much to listen, we gather the actual content, which lives in $V$ . Splitting "who to listen to" (K) from "what you hear" (V) lets the model tune them independently.

PICTURE. Token 3's weight bar times each value row, summed, gives one new row: token 3's context-aware output.

That completes one attention head — exactly Self-Attention. Now the "multi".

Step 6 — Split into $h$ heads: run the same trick on narrow slices

WHAT. Instead of one head using all $d_{model}$ columns, we cut the projections so each of $h$ heads works in a narrow slice of width $d_{k} = \frac{d _{model}}{h} .$ Every head $i$ gets its own $W_{i}^{Q}, W_{i}^{K}, W_{i}^{V}$ and runs Steps 3–5 by itself.

WHY narrow slices? So the total work stays the same. Attention cost grows like $n^{2} d$ . One fat head of width $512$ costs $n^{2} \cdot 512$ . Eight thin heads of width $64$ cost $8 \times n^{2} \cdot 64 = n^{2} \cdot 512$ — identical total, but now 8 different viewpoints. One head can chase grammar, another meaning, another position.

WHY different $W_{i}$ per head? If every head shared one projection they'd compute identical grids — 8 copies of the same answer, zero gain. Different reshapers force different questions.

PICTURE. The wide grid is sliced into $h$ colored bands; each band is an independent little attention machine.

Degenerate case $h = 1$ . Then $d_{k} = d_{model}$ and multi-head collapses back to plain single-head Self-Attention — the parent formula still holds. Failure case $h ∤ d_{model}$ . If $h$ does not divide $d_{model}$ evenly, $d_{k}$ isn't a whole number and the split is illegal — real implementations require $d_{model}$ be a multiple of $h$ .

Step 7 — Concatenate, then one final mix $W^{O}$

WHAT. Lay the $h$ head-outputs side by side (each $n \times d_{v}$ ) to rebuild a full $n \times d_{model}$ grid, then multiply by a last learned matrix $W^{O}$ : $MultiHead = Concat (head_{1}, \dots, head_{h}) W^{O}$

WHY concatenate and not average? Averaging would blend the 8 viewpoints into mush, erasing the specialization we paid for. Concatenation stacks them side by side, keeping every head's answer intact.

WHY still need $W^{O}$ ? After concatenation the heads sit in separate lanes that have never talked. $W^{O}$ is the mixer that lets grammar-info and meaning-info combine, and it snaps the width back to $d_{model}$ so the next layer (Transformer Encoder Layer) can consume it.

PICTURE. Colored head-blocks slotted together into one wide grid, then squeezed through $W^{O}$ into the final output.

The one-picture summary

Everything above on a single canvas: one sentence-grid → split into colored heads → each head does dot-product + scale + softmax + value-blend → concatenate → $W^{O}$ → output.

Recall Feynman retelling — say it back in plain words

A sentence walks in as a grid of numbers, one row per word. We make three copies reshaped by learned matrices: questions (Q), labels (K), and contents (V). To decide how much word A cares about word B, we dot A's question with B's label — big dot means "yes, that answers me". We shrink the numbers by $d_{k}$ so they don't blow up, then run softmax so each word's caring-amounts add to 1. Each word then grabs everyone's content, weighted by how much it cared — that's one head's answer. Now the twist: we don't do this once, we do it $h$ times in skinny slices so the total cost is unchanged, and each slice gets its own reshapers so it hunts a different pattern (grammar, meaning, position). Finally we lay the skinny answers side by side and pass them through one last matrix $W^{O}$ that lets the different patterns talk and resizes everything back. Out comes a grid the same shape as we started with — but every word now knows about its neighbors.

Recall Check yourself

Why divide scores by $d_{k}$ ? ::: Long rows make dot products grow large; dividing by $d_{k}$ rescales them so softmax stays sensitive and gradients survive. Why concatenate heads instead of averaging? ::: To keep each head's distinct pattern; $W^{O}$ then mixes them deliberately. What is $d_{k}$ when $d_{model} = 512, h = 8$ ? ::: $512/8 = 64$ . What happens to multi-head attention when $h = 1$ ? ::: It reduces to ordinary single-head self-attention.

Related views: Cross-Attention, Positional Encoding, Attention Visualization.