Attention heads and the causal mask

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

For each head, predict where position 9 will look before you tap it.
Turn off the causal mask and look at which heads change. The previous-token and first-token heads barely change: their high scores are never on later tokens, so the future gets only a sliver of weight.

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