Attention in Transformers

One Character at a Time

The brass type slugs in front are what the model has read: “the queen. the qu”. The glass tubes at the back are its guess for the next character, the height of each column its probability. It says e, with more than 90 % confidence.

This model knows 30 characters: 26 letters, space, full stop, comma and apostrophe. It was trained on six public-domain novels and is made of 27,166 numbers. GPT-3 has 175 billion, in 96 layers of 96 heads each. The job is the same: look back, guess the next.

2 layers × 4 heads · 32 numbers per character · 27,166 parameters · running live on this page

Characters Become Numbers

The model knows no letters, only numbers. Each character is looked up in a table and becomes 32 numbers, its embedding. The same letter gets the same 32 numbers wherever it stands.

But where it stands matters too. The 2017 Transformer paper adds to each position a set of sines and cosines, from fast to slow. The first stretch of each glass tube is the two added together: its colours are the 32 numbers, copper for positive, blue for negative.

Embedding (one per character) + position (one set of sines and cosines per place) = the 32 numbers that enter layer 1

Queries and Keys

From those 32 numbers every head makes three things of 8 numbers each: a query (what am I looking for), a key (what am I) and a value (what I hand over if chosen).

The plate lifted now is layer 1, head 4. The query of the u on the far right is multiplied, number by number, with the key of every character before it, and divided by √8. Each small dial’s needle is that score: the higher it reads, the more that character looks like what it is after. For this head, the highest is the q just before.

score = query · key / √8 · needles computed live on this page

Only Looking Back

The scores go through a softmax: take the exponential of each and divide by their sum. The results are all positive and add up to exactly 1. These are the attention weights: the thicker and brighter the thread, the larger the weight.

Each character may look only at itself and what came before, never ahead: in training the model must guess exactly what comes after, and looking would be cheating. This is the causal mask; the slanted shutter on the machine blocks everything to the right.

softmax(scores) → weights, each row adds up to 1 · the weight from i to any j > i is always 0

Bringing Back What It Found

Each head adds up the characters’ values, each times its weight, and sends the sum back into the query character’s tube. The beads running along the threads are those values; a bead’s size is its weight.

What comes back is not swapped in but added. Each character’s tube is one channel from front to back, and every layer only adds a little to it. A 2021 paper calls this channel the residual stream.

head output = Σ weight × value, added back to the residual stream

Cut Open

Front to back: type slugs, embedding, layer 1’s four attention heads, layer 1’s feed-forward network (MLP), layer 2’s four heads, layer 2’s MLP, a last normalisation, the read-out.

The MLP works on each character on its own, looking at no other: its 32 numbers become 128, pass through a bent function, become 32 again and are added back to the tube. The row of fine needles in the cut-open housing is the 128 neurons for the chosen character.

At the end, the read-out turns the tube’s last 32 numbers into 30 scores, and a softmax turns those into the probabilities in the glass tubes at the back.

embedding → [attention → MLP] × 2 → read-out · every step adds to the same residual stream

Eight Heads, Eight Jobs

The four heads of a layer read the same numbers, yet each has learned something different. Nobody assigned the jobs; training divided them.

In this model, layer 1’s heads 2 and 4 put more than half their attention on the character just before; head 1 looks two back more than anywhere else; head 3 leans towards spaces, the gaps between words. Layer 2’s heads 3 and 4 do something stranger, in the next chapter.

Layer 1: previous character (heads 2, 4) · two back (1) · spaces (3)

Copying What Came Next Last Time

Read this nonsense: “qzxv kplw qzxv kp”. The model never met these letters together in any book, yet it guesses l next, at about 80 %.

It takes both layers. Layer 1’s previous-character heads first mark each character with “who came before me”. A layer 2 head then takes the current p and asks: whose previous character was p? It finds the l from the first time round and brings it over. A B … A → B. Researchers at Anthropic named these induction heads in 2021. When there is nothing to copy, these two heads mostly rest on spaces.

Layer 2, heads 3 and 4: from the second p to the l after the first p

It Appears Suddenly

The drum on the left records the whole training run: 9,000 steps, each reading 48 stretches of 64 characters. Turn it, and the machine takes on the weights of that moment.

In the first 1,000 steps the model learned common pairs of letters but could not copy at all: on unseen repeated nonsense it got the second time round right only about one time in ten. In the next 500 steps the induction heads suddenly formed, and that rose to nearly 80 %. A 2022 study saw such a sudden change in many models of different sizes, often leaving a small bump in the training loss.

Step 1,000 → step 1,500: the repeated part right 14 % → 78 % of the time

Pulling a Head Out

Setting a head’s output to zero is like pulling it out of the rack. On 200 unseen stretches of repeated nonsense, this model gets 96 % of the characters right the second time round.

Pull out layer 2, head 4, and that falls to 33 %. Pull out the head it relies on, layer 1, head 4 (previous character), and it falls to 28 %. Pull out layer 2, head 1, and almost nothing changes. Whether a head really matters, pulling it once tells.

Repeated nonsense, second time right: all heads 96 % · without L2H4 33 % · without L1H4 28 % · without L2H1 95 %

The Square and the Cache

To read n characters, each head computes about n²/2 scores: twice the characters, four times the scores. GPT-3 reads 2,048 tokens at a time: in its full-attention layers, over two million scores per head. This is the quadratic cost of attention.

When a model writes one token at a time, the keys and values of everything before do not change. They can be kept, and only the new token needs a query. This is the KV cache. It costs memory: at GPT-3’s size (every layer keeping every token), about 2.36 million numbers per token, about 4.7 MB in 16-bit floats; 2,048 tokens take about 9.7 GB.

With the cache: ≈ n dot products per new token · without: ≈ n²/2