A single attention head·6 min

What attention computes

0 / 1

Attention is the core of a transformer. Here is the whole idea in one sentence: each token writes a short summary of the tokens it's allowed to look back at. The query is what this token is looking for; the keys are what every token advertises about itself; the value is what each token actually contributes to the summary. Attention is a content-addressed, weighted lookup — and over the next five steps you'll build one causal scaled-dot-product head, the exact unit GPT-2 runs 12 times in parallel (Chapter 6).

How a head works

You're given three matrices q, k, v, each shape [T, d] (T positions, head dimension d). For each query position i:

Scores — dot every query against every key, scaled by 1/√d. The 1/√d keeps the dot products from growing with d and saturating softmax.
Causal mask — a position may only attend to itself and earlier positions. Set scores[i][j] to a huge negative number for j > i (so softmax sends those to 0).
Softmax over the last axis turns each row into weights that sum to 1.
Output — the weighted sum of value vectors.

Data flow, with shapes

q, k, v : [T, d] each

   q · kᵀ /√d        mask         softmax (rows)        weights · v
[T,d]×[T,d] ─▶ scores ─▶ masked ─▶  weights  ─────────────▶  out
              [T,T]      [T,T]       [T,T]                    [T,d]

Three shapes to keep straight

q, k, v are [T, d]. scores and weights are the square [T, T] — "who attends to whom." out is back to [T, d]. If a shape isn't one of those three, something is off.

Read a tensor Scaled scores