Chapter 4 · The transformerPuff 16 of 23

Many heads looking at once

“Puff found a shiny key and kept it.” To understand “it”, the model needs several things at once: which word came just before, who did the keeping, and what “it” stands for. One attention step gives each word one set of percentages. So how does it follow several links at the same time?

Guess first

Transformers run many attention “heads” side by side instead of one. Why?

Play with it

Fig. 1 · Multi-head attention Three heads, three kinds of link

Toy model: three heads, each with its own hand-made 2-number queries and keys. Head 1’s lists come from each word’s position, heads 2 and 3 from word features. Real heads learn their lists from data and are messier. In each grid, a row is a word looking back and the columns are the words it can see, in sentence order.

Tap a word to see its row in every head. Then look at one head at a time.

How it works

Tap a step to see it in the figure.

Check yourself

Researchers find a head that always looks at the word just before. What is it probably useful for?

Go deeper: the mathThe multi-head formula, head sizes, and what real heads do

The formula

Each head i has its own three learned tables, W_i^Q, W_i^K and W_i^V, and does ordinary attention:

head_i =
Attention(XW_i^Q, XW_i^K, XW_i^V)

The heads’ outputs are joined end to end (concatenated) and mixed by one more learned table, W^O:

MultiHead(X) =
Concat(head₁, …, head_h) W^O

How big is each head?

The word’s list of d_model numbers is shared out between the h heads:

d_head = d_modelh

This page: 6 ÷ 3 = 2 numbers per head.
The original Transformer (base model): 512 ÷ 8 = 64.
GPT-2’s smallest model: 768 ÷ 12 = 64.

Joined back up, the heads give h × d_head = d_model numbers again, so the next step sees a list the same length as before. Because each head is smaller, many heads cost about the same as one full-size head.

This page’s heads, live

Shares for “it” in each head (each column adds to 100%):

Head 1 builds each word’s query and key from its position, like a clock hand that turns 0.7 radians per word: k_j = 3.5 · [cos 0.7j, sin 0.7j], and word i’s query points where the previous word’s hand points, q_i = 3.5 · [cos 0.7(i − 1), sin 0.7(i − 1)]. Their dot product is 12.25 · cos 0.7(i − 1 − j), biggest when j = i − 1. Heads 2 and 3 use two hand-made features each; every other list is zeros, so a word with nothing to look for spreads its attention evenly:

What real heads do

Clark and colleagues (2019) looked inside BERT, a well-known transformer, and found heads that attend to the previous or next word, heads that link objects to their verbs and a “the” to its noun, and heads that link a pronoun to what it refers to. Many heads also park their attention on one special token when they have nothing to do. Voita and colleagues (2019) found that in a translation model a few specialised heads (tracking position, grammar or rare words) do most of the work, and many others can be removed with little loss. Olsson and colleagues (2022) showed previous-word heads teaming up with “induction heads” to copy patterns that appeared earlier in the text.

Real heads are learned, not designed, and many don’t have one clean job. The three here are hand-built to show the kinds of pattern researchers have found.

Sources

A. Vaswani et al., “Attention Is All You Need”, 2017. Section 3.2.2 (multi-head attention; 8 heads, d_model = 512).
K. Clark, U. Khandelwal, O. Levy and C. D. Manning, “What Does BERT Look At? An Analysis of BERT’s Attention”, 2019.
E. Voita, D. Talbot, F. Moiseev, R. Sennrich and I. Titov, “Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned”, 2019.
C. Olsson et al., “In-context Learning and Induction Heads”, Anthropic, 2022. GPT-2 sizes: A. Radford et al., “Language Models are Unsupervised Multitask Learners”, OpenAI, 2019, and the model settings in OpenAI’s gpt-2 code release. The toy heads on this page are hand-made.