The attention mechanism

Embeddings give every token a starting vector, but that vector alone can't know anything about context — the word "bank" needs a different meaning next to "river" than next to "loan". Attention is the mechanism that lets every token pull in information from every other token before deciding what it means.

Query, key, value #

Every token produces three vectors from its embedding, via three learned weight matrices:

• Query ($Q$) — what this token is looking for
• Key ($K$) — what this token offers, for others to match against
• Value ($V$) — the actual information this token contributes if selected

A token's query is compared against every other token's key to produce a relevance score, which is then used to weight that token's value:

The $\sqrt{d_k}$ term just keeps the dot products from growing too large as the vector dimension increases — without it, softmax would saturate and gradients would vanish during training.

Two things to notice while playing with the grid. First, each row sums to exactly 1 — softmax always produces a probability distribution, so a token's attention is a budget it spends across the other tokens. Second, toggle the $\sqrt{d_k}$ scaling off: the same scores produce much sharper weights, which is exactly the saturation problem the scaling prevents during training.

Why multiple heads #

A single attention computation can only learn one notion of "relevance." Real models run many attention computations in parallel — heads — each with its own $Q$, $K$, $V$ projections, so the model can track several kinds of relationships at once (e.g. one head tracking grammatical subject-verb agreement, another tracking coreference across a paragraph).

# simplified single-head attention
def attention(Q, K, V):
    scores = Q @ K.T / sqrt(d_k)
    weights = softmax(scores, axis=-1)
    return weights @ V

The dimension arithmetic is worth internalizing, because it explains why heads are nearly free:

The model's dimension is split between heads, not duplicated: 12 heads at $d_{\text{model}} = 768$ each operate on 64-dimensional sub-vectors, and their outputs are concatenated back to 768 before the output projection. Total QKV parameters stay the same as one big head would cost — you're trading "one high-capacity relevance map" for "12 specialists at lower resolution each." Empirically the specialists win — with one honest caveat: per-head ablations show many heads are redundant or prunable with negligible loss; only a handful per layer are load-bearing, but those carry relationships (coreference, positional copying, induction — a term the interpretability chapter defines) that no single high-capacity head provides. Head-level patterns were readable with simple probing years before heavyweight mechanistic-interpretability tooling existed.

Switch between four specialists on the same sentence:

This lesson has exercises attached — working through the score computation by hand for a 3-token sequence — once the exercises layer ships.

Next: the transformer block — the repeatable unit these attention heads live inside.