Transformers, Part 2: Self-Attention & Multi-Head Attention
How self-attention works — softmax(QKᵀ/√dₖ)V computed on a real sentence, why we scale by √dₖ, and how multi-head attention splits d_model across heads.
This article has an interactive companion
Don't just read it — step through it interactively, one state change at a time.
After tokenization and embeddings, every word in a sentence is a vector — but each vector is context-free. The embedding for "bank" is identical whether the sentence is about a river or a savings account. That is the problem self-attention solves: it lets every token look at every other token and rewrite itself into a context-aware vector.
Self-attention is the operation that made transformers work, and it is far simpler than the notation suggests. Strip away the matrices and it is a weighted average — a soft lookup where a token asks "which other tokens matter to me?" and blends their information in proportion to the answer.
This walks through the exact mechanism the YeetCode deep dive computes live: the scaled dot-product formula on a real three-word sentence, why the √dₖ divisor is not optional, and how multi-head attention runs several of these lookups in parallel. It stops where attention stops — residual connections, layer norm, and the feed-forward block are the next article.
Self-attention is a soft dictionary lookup
A normal dictionary is a hard lookup: you give a key, you get back exactly one value. Self-attention softens that. Instead of matching one key, each token compares itself against every key, turns those comparisons into weights that sum to 1, and returns a blend of all the values weighted by relevance.
Concretely, for a token like "cat" in "the cat sat", the mechanism:
- Scores "cat" against every token in the sentence (including itself).
- Converts those raw scores into a probability distribution with softmax.
- Returns a weighted sum of every token's value vector, using those probabilities as the mixing weights.
The output is a new vector for "cat" that has absorbed information from its neighbours. Do this for all tokens in parallel and you have one attention layer. Because every token attends to every other token, the operation is O(n²) in sequence length — that quadratic cost is the reason long context windows are expensive.
Query, Key, and Value: three projections of one token
Attention never scores raw embeddings directly. Each token's embedding x is first projected through three separate learned weight matrices into three roles:
Q = x · Wqᵀ "what I'm looking for"
K = x · Wkᵀ "what I offer to others"
V = x · Wvᵀ "what I'll pass on if chosen"The query and key live in the same space so they can be compared with a dot product — a large Q·K means "this key is relevant to this query." The value is separate: it holds the actual information that gets mixed into the output. Splitting "how relevant am I?" (K) from "what do I contribute?" (V) is what gives attention its flexibility. A token can be highly attended-to while contributing a value quite different from its key.
All three are learned. In an untrained model — like the one the interactive uses so you can inspect real arithmetic — Wq, Wk, and Wv are random, so the attention pattern is real but not yet meaningful. Learned coreference (knowing "it" refers to "the animal") only emerges after training.
The formula: softmax(QKᵀ/√dₖ)V
The whole operation is one line:
Attention(Q, K, V) = softmax( Q Kᵀ / √dₖ ) VRead left to right:
Q Kᵀ— every query dotted against every key, producing an n×n matrix of raw similarity scores./ √dₖ— divide by the square root of the key dimension to keep scores in a sane range (the next section explains why).softmax(...)— normalize each row into a probability distribution over keys, summing to exactly 1.... V— multiply the weights by the value vectors and sum, producing the context vector.
function selfAttention(X, q) { // X: position-aware embeddings
const Q = X.map((x) => matVec(Wq, x)); // queries (n × d_k)
const K = X.map((x) => matVec(Wk, x)); // keys (n × d_k)
const V = X.map((x) => matVec(Wv, x)); // values (n × d_k)
const raw = K.map((k) => dot(Q[q], k)); // Qq · Kj for all j
const scaled = raw.map((s) => s / Math.sqrt(d_k));
const w = softmax(scaled); // Σ = 1
const out = V[0].map((_, d) =>
w.reduce((s, wj, j) => s + wj * V[j][d], 0)); // weighted sum of values
return out; // the query's new representation
}That out vector replaces the token's input embedding as its context-aware meaning.
Walkthrough: attention for "cat" in "the cat sat"
Take the sentence the cat sat with d_model = 24, a single head, and √dₖ scaling on. We resolve the query token "cat" (index 1). Every number below is the real output of the engine, not an illustration.
| Stage | Operation | Result for query "cat" |
|---|---|---|
| Embed | embed("cat") + PE(1) | one 24-dim position-aware vector |
| Project | Q = xWqᵀ, K = xWkᵀ, V = xWvᵀ | Q, K, V, each 24-dim |
Score QKᵀ | dot(Q_cat, K_j) for each key | the → 3.877 · cat → 14.905 · sat → 8.293 |
Scale /√24 | divide each by 4.899 | the → 0.791 · cat → 3.042 · sat → 1.693 |
| Softmax | normalize to a distribution | the → 7.7% · cat → 73.3% · sat → 19.0% (Σ = 1.000) |
| Weight V | Σ wⱼ · Vⱼ | context ≈ [-2.26, 2.17, -1.41, -1.52, …] |
The raw score for "cat" against itself (14.905) is far larger than against "the" (3.877), so after softmax "cat" keeps 73.3% of its own value and mixes in 19% of "sat" and 7.7% of "the." The final context vector is that blend.
One honest caveat the deep dive makes explicit: because these weights are untrained and random, "cat" attending mostly to itself is an artifact of the initialization, not linguistic understanding. The mechanism is exactly what a trained GPT runs; only the learned patterns are missing.
Why the √dₖ scaling exists
That /√dₖ divisor looks like a fudge factor. It is load-bearing. Dot products grow with dimension: add more components and Q·K drifts to larger magnitudes. Feed large scores into softmax and it saturates — one weight rockets toward 1.0 and the rest collapse to nearly 0. A near one-hot distribution kills the gradient flowing to every non-winning token, so the model can barely learn.
Watch what happens to the exact same sentence with scaling turned off (dividing by 1 instead of 4.899):
| the | cat | sat | |
|---|---|---|---|
| Scaled by √dₖ | 7.7% | 73.3% | 19.0% |
| Unscaled | 0.0% | 99.9% | 0.1% |
Without scaling the softmax spikes to 99.9% on a single token — the "soft" lookup has hardened into a "hard" one, and gradients through the other tokens vanish. Dividing by √dₖ keeps scores in a range where softmax stays soft and trainable. That is the entire reason the term is in the formula.
Multi-head attention: many perspectives in parallel
A single attention head learns one kind of relationship. Language has many at once — subject–verb agreement, coreference, local word order, long-range dependency. Rather than force one head to capture all of them, transformers run h heads in parallel, each with its own Wq, Wk, Wv.
The trick is that heads don't add width — they split it. The model dimension is divided across heads:
d_k = d_model / hWith d_model = 24 and h = 4, each head works in a 6-dimensional subspace (24 / 4 = 6). GPT-2 uses 768 / 12 = 64. The division must be exact: if d_model isn't divisible by h, d_k is fractional and the architecture is invalid — the interactive turns the dimension card red to flag exactly this.
Each of the four heads runs the identical softmax(QKᵀ/√dₖ)V, but on different projections, so each produces a different soft distribution over the sentence. Then the outputs are stitched back together:
function multiHead(X, q, h) {
const d_k = d_model / h; // split the model dim across heads
const heads = [];
for (let i = 0; i < h; i++) { // each head: its OWN Wq/Wk/Wv
heads.push(attention(X, q, Wq[i], Wk[i], Wv[i], d_k));
}
const concat = heads.flatMap((head) => head.out); // length h·d_k
const out = matVec(Wo, concat); // project h·d_k → d_model
return out; // fused multi-perspective vector
}Two steps finish the job. Concatenate the four 6-dim head outputs into one 24-dim vector (4 × 6 = 24, exactly d_model again). Then project it through the output matrix W_O, which maps h·dₖ → d_model. W_O is not a reshape — it lets the heads' perspectives interact and mix, fusing "several parallel lookups" back into one representation per token, ready for the next layer.
What to remember
- Self-attention is a soft lookup: score a query against all keys, softmax to weights that sum to 1, return a weighted blend of values.
- Q, K, V are three learned projections of the same embedding — query and key negotiate relevance, value carries the payload.
- The formula is
softmax(QKᵀ/√dₖ)V, and the√dₖterm is what stops softmax from saturating into a one-hot spike. - Multi-head splits
d_modelacrosshheads (d_k = d_model/h, exact division required), runs attention per head, then concat +W_Ofuses them. - Attention is O(n²) in sequence length — every token attends to every other token.
This is the end of the attention story. What sits around attention — residual connections, layer normalization, and the feed-forward network that together form a transformer block — is the subject of the modern-architecture article below.
Keep going
Step through every number yourself in the interactive deep dive: type your own sentence, pick a query token, watch real Q·K scores become softmax weights, toggle √dₖ scaling to see the saturation, and slide the head count to watch d_k = d_model/h recompute live.
Then continue the series:
- What Changed Since 2017: RMSNorm, SwiGLU, RoPE, KV Cache & GQA — how residuals, norm, and the FFN evolved into today's blocks.
- Train Your First Language Model: From Raw Text to Generation — put attention inside a full training loop.
- LLM Fine-Tuning, Part 1: When to Fine-Tune & How LoRA Works — adapt a pretrained transformer cheaply.
- Supervised Fine-Tuning: How a Base Model Becomes an Assistant — turn a next-token predictor into a chat model.
FAQ
What does softmax(QKᵀ/√dₖ)V actually compute?
It computes a context-aware vector for each token. QKᵀ scores every query against every key with dot products, /√dₖ scales those scores so softmax doesn't saturate, softmax turns each row into weights that sum to 1, and multiplying by V returns a weighted average of the value vectors. The result is a new representation for each token that blends in information from every other token, weighted by relevance.
Why divide the attention scores by √dₖ?
Dot products grow in magnitude as the key dimension dₖ increases. Large scores push softmax toward a one-hot distribution where a single token gets nearly all the weight and every other weight collapses to almost zero. That saturation starves the non-winning tokens of gradient and makes training unstable. Dividing by √dₖ rescales the scores back into a range where softmax stays soft and differentiable — in the worked example it changed a 99.9% spike into a healthier 73.3% peak.
What is the difference between Query, Key, and Value?
They are three separate learned projections of the same input embedding. The query represents what a token is looking for, the key represents what a token offers, and the value represents the information a token contributes if it gets attended to. Query and key are compared with a dot product to produce relevance scores, while the value is the payload that actually gets mixed into the output. Separating relevance (K) from content (V) is what lets a token be highly relevant while contributing something distinct.
How does multi-head attention differ from single-head attention?
Multi-head attention runs h independent attention operations in parallel instead of one. It splits the model dimension so each head works in a d_k = d_model / h subspace with its own Q/K/V projections, letting different heads capture different relationships — syntax, coreference, locality. Their outputs are concatenated back to d_model and projected through W_O. Single-head attention is just the h = 1 case where one head uses the full d_model.
Does this cover the whole transformer block?
No. Self-attention and multi-head attention are the core, but a full transformer block also wraps them in residual connections, layer normalization, and a position-wise feed-forward network. Those components stabilize training and add non-linear capacity on top of attention. They are covered in the modern-architecture article rather than here — this piece deliberately stops at multi-head attention.
Make it stick: run this one yourself
Don't just read it — step through it interactively, one state change at a time.