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%

Challenges

0 of 3 complete

On the Induction head with the mask on, find a query other than the second cat that looks mostly at the very word that comes next. (not yet)
On the Same word head with the mask on, pick a token that has no earlier copy of itself and see where it looks. (not yet)
Turn off the causal mask and find a query whose favorite token comes after it. (not yet)

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.

A masked attention head in a few linesPython with NumPy

import numpy as np

tokens = ["the", "cat", "chased", "the", "dog", ".", "the", "cat", "chased", "the"]
n = len(tokens)

# The demo's induction head scores pairs by rule: 5 where the key's previous
# token matches the query token, 2 for position 0 (a fallback), 0 elsewhere.
scores = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        if j >= 1 and j - 1 != i and tokens[j - 1] == tokens[i]:
            scores[i, j] = 5
        elif j == 0:
            scores[i, j] = 2

# Causal mask: no query may see a key that comes after it.
future = np.triu(np.ones((n, n), dtype=bool), k=1)
scores[future] = -np.inf

weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)

top = weights[7].argmax()
print(f"cat (7) looks most at {tokens[top]} ({top}): {weights[7, top]:.0%}")
print("total weight on future tokens:", weights[future].sum())

# Splitting into heads: GPT-2 small has d_model = 768 and 12 heads.
d_model, h = 768, 12
print("per-head width:", d_model // h, "| projection weights:", 4 * d_model**2)

cat (7) looks most at chased (2): 92%
total weight on future tokens: 0.0
per-head width: 64 | projection weights: 2359296

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