Many heads and the causal mask

In the last lesson a token’s query produced one set of weights, and the token received one weighted average of the values. That is a tight budget. A word may need several kinds of information at once: which noun it refers to, what word came just before it, whether it appeared earlier in the text. With one softmax, all of those needs compete for the same total weight of 1.

Transformers handle this by running several attention operations in parallel, called heads, and they make text generation possible with one more ingredient, a mask that stops tokens from seeing the future.

Multi-head attention

Multi-head attentionRunning several attention heads in parallel, each with its own query, key, and value projections, then concatenating their outputs and mixing them with one more matrix.Open in glossary runs $h$ separate Attention headOne independent attention computation inside a multi-head attention layer, with its own learned projections and its own pattern of weights.Open in glossary. Each head has its own learned projections, so it computes its own queries, keys, and values and produces its own pattern of weights:

\mathrm{head}_i = \mathrm{Attention}(X W^Q_i,\; X W^K_i,\; X W^V_i)

\mathrm{MultiHead}(X) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_h)\, W^O

The heads do not add width. Each one works in $d_k = d_{\text{model}} / h$ dimensions, and concatenating $h$ of them gives back a vector of width $d_{\text{model}}$ , which the output matrix $W^O$ mixes together. The original transformer used $d_{\text{model}} = 512$ and $h = 8$ , so each head had 64 dimensions.

Splitting the model width into heads

More heads means narrower heads. The total size of the attention layer does not change.

One token's 512-dimensional vector, cut into 8 heads of 64 dimensions. Each head has its own query, key, and value projections of shape 512 × 64, runs attention independently, and the 8 outputs are concatenated back to width 512 and multiplied by W_O (512 × 512).

Model

Transformer baseGPT-2 smallGPT-3 175B

Number of heads, h8

Model width, d_model512

Per-head width, d_k = d_model / h64

Projection weights, 4 × d_model²1,048,576

Try this

Pick GPT-2 small and move the heads slider. The per-head width changes; the number of projection weights stays at 2,359,296.
Pick GPT-3 175B: 96 heads of 128 dimensions each, inside a 12,288-dimensional model.

In standard multi-head attention, where each head has width $d_{\text{model}}/h$ , the per-head projections together form $d_{\text{model}} \times d_{\text{model}}$ matrices, so one attention layer has $4\,d_{\text{model}}^2$ projection weights (plus biases, if the model uses them), whatever $h$ is. Some newer models share keys and values across heads to save memory, which changes this count; the KV cache lesson covers that. The number of heads is a choice about how to divide that capacity: many narrow heads, each free to attend somewhere different, or a few wide ones.

Hiding the future: the causal mask

A language model is trained to predict the next token at every position of a training sequence at the same time. That parallelism is a big part of why transformers train efficiently. But it creates a problem: if position 5 can attend to position 6, it can simply read the answer it is supposed to predict.

The fix is the Causal maskA mask that stops each token from attending to tokens after it, by setting those attention scores to negative infinity before the softmax. Used by models that generate text left to right.Open in glossary. Before the softmax, every score where the key comes after the query is replaced by $-\infty$ :

\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right), \qquad M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases}

Since $e^{-\infty} = 0$ , those positions get exactly zero weight, and the weight matrix becomes lower triangular. Each token sees itself and everything before it, never after. This matches how generation works: when the model produces token 6, tokens 7 onward do not exist yet.

Not every transformer uses the mask. Encoder models such as BERT let every token see the whole input in both directions, which suits tasks like classifying or searching text, but they are not built to generate text one token at a time.

What heads learn to do

Nobody tells a head what to attend to. Its pattern emerges from training. Interpretability research has nonetheless found patterns that show up again and again. The demo builds five of them on one short sentence so you can compare them side by side.

Five heads, one sentence

Each head applies its own scoring rule to the same tokens. Pick a head, then pick a query token to see where it looks.

Each token finds an earlier copy of itself and looks at the token that came next, which helps predict a repeat.

Induction heads are described by Elhage et al. (2021) and Olsson et al. (2022), who tie them to in-context learning.

Attention weights for the Induction head. Rows are queries, columns are keys. Cells above the diagonal are masked.
the	cat	chased	the	dog	.	the	cat	chased	the
1.00	masked	masked	masked	masked	masked	masked	masked	masked	masked
0.88	0.12	masked	masked	masked	masked	masked	masked	masked	masked
0.79	0.11	0.11	masked	masked	masked	masked	masked	masked	masked
0.05	0.94	0.01	0.01	masked	masked	masked	masked	masked	masked
0.65	0.09	0.09	0.09	0.09	masked	masked	masked	masked	masked
0.60	0.08	0.08	0.08	0.08	0.08	masked	masked	masked	masked
0.02	0.48	0.00	0.00	0.48	0.00	0.00	masked	masked	masked
0.05	0.01	0.92	0.01	0.01	0.01	0.01	0.01	masked	masked
0.05	0.01	0.01	0.91	0.01	0.01	0.01	0.01	0.01	masked
0.02	0.32	0.00	0.00	0.32	0.00	0.00	0.32	0.00	0.00

Causal mask (each token sees only itself and earlier tokens)

HeadInduction

Querycat (7)

Looks most atchased (2), 92%

These heads are constructed: each one scores token pairs with a simple rule, then the real mask and softmax turn scores into weights. The rules imitate patterns found in trained models; a trained head gets its scores from learned queries and keys and is rarely this clean.

Try this

Choose Induction and tap the second cat (position 7). It looks at chased, the word that came after “cat” the first time, and “chased” is exactly what comes next.
Still on Induction, tap the last the (position 9). Its attention splits between cat and dog, the two different words that followed earlier copies of “the”.
Choose Same word, tap the first the, and switch the causal mask off. Now it can see later copies of itself. In a model that generates text, those later tokens would not exist yet.
Choose First token and compare the mini maps. Every row piles onto position 0.

What these patterns are, and where they were described:

Previous-token heads attend to the token just before the current one. They are simple, common, and useful as a building block for other heads.
Duplicate-token heads attend to earlier copies of the current token. Wang et al. (2022) found them in GPT-2 small as part of the circuit it uses to complete sentences like “When Mary and John went to the store, John gave a drink to ___”.
Induction headAn attention head that, at a token A, looks back for earlier copies of A and attends to the token that followed them, helping the model continue repeated patterns.Open in glossary complete patterns of the form A B … A → B. Elhage et al. (2021) and Olsson et al. (2022) described them and linked them to in-context learning, a model’s ability to pick up patterns from its prompt.
First-token heads put most of their weight on the first position whenever nothing else is relevant. Xiao et al. (2023) studied this as “attention sinks” and found it matters for running models on long streams of text. In trained models this is less a separate kind of head than a habit shared by most heads beyond the first couple of layers; the demo gives it a head of its own to make it visible.
Broad heads spread attention evenly, averaging over the context. Clark et al. (2019) observed broad heads in BERT’s lower layers.

How an induction head works with two layersOptional

A single head cannot implement the induction rule alone, because a key only describes its own token. The trick, described by Elhage et al. (2021), takes two layers:

In an earlier layer, a previous-token head copies information about token $j - 1$ into position $j$ ‘s vector. Now each position “knows” what came before it.
In a later layer, the induction head’s query at the current token $A$ looks for keys whose previous token was $A$ . That is the position holding $B$ , and its value carries $B$ forward to help predict the next token.

In the demo the rule is applied directly, so one head shows the end result. The causal mask also makes this efficient during generation: since earlier tokens never attend to later ones, their keys and values never change as new tokens arrive, so a model can compute them once and reuse them. That reuse is the KV cache, covered later in this track.

Key ideas

Multi-head attention runs $h$ attention heads in parallel, each in $d_k = d_{\text{model}}/h$ dimensions, then concatenates them and mixes with $W^O$ .
Splitting into heads does not change the number of projection weights, $4\,d_{\text{model}}^2$ per layer. It lets the model attend to several things at once.
The causal mask sets scores for future positions to $-\infty$ so their weights are exactly 0. It makes parallel training honest and matches left-to-right generation.
Trained heads show recurring patterns, including previous-token, duplicate-token, induction, and first-token heads, but many heads are messier than these examples.