Attention and Transformers

The question. How can a machine translate a fifty-word sentence if only a handful of numbers may pass from the part that reads to the part that writes?

That was the trouble with the translator of Unit 17. Its encoder reads "I drink tea" one word at a time. After each word it writes a note — a short list of numbers: $\mathbf h_1$ after "I", $\mathbf h_2$ after "drink", $\mathbf h_3$ after "tea". Then it hands the decoder only the last note, the summary. A five-word sentence and a fifty-word sentence must squeeze through the same few numbers. That squeeze is the bottleneck. It is the first interpreter, speaking from memory.

The fix (Bahdanau, Cho and Bengio, 2015) is the second interpreter. Keep every note. Before the decoder writes each word, it makes three small moves.

1. Score every note. How well does this note match what I am about to write? The simplest score is the dot product (Unit 3) of the decoder's current state $\mathbf q$ with the note.
2. Turn the scores into shares with softmax (Unit 14). The shares are positive and add up to 1. They are called the attention weights.
3. Blend the notes with those shares. The blend is the context vector. The decoder uses it, together with its own state, to choose the next word.

By hand. We keep the numbers tiny and make them by hand, so that you can see everything. Each note simply points at its own word: $\mathbf h_{\text{I}}=(1,0,0)$, $\mathbf h_{\text{drink}}=(0,1,0)$, $\mathbf h_{\text{tea}}=(0,0,1)$. (Real notes are long lists that mix many words together.) The decoder has written मैं (main, "I") and is about to write चाय (chai, "tea"). Its state — the decoder's own running note, as in Unit 17 — is, say, $\mathbf q=(0,\ 0.5,\ 2)$.

$$\begin{aligned}\text{scores: }&\mathbf q\cdot\mathbf h=(0,\ 0.5,\ 2)\\ \text{softmax: }&\frac{(e^{0},\ e^{0.5},\ e^{2})}{1+1.649+7.389}\\ &\approx(0.100,\ 0.164,\ 0.736)\\ \text{context: }&0.100\,\mathbf h_{\text{I}}+0.164\,\mathbf h_{\text{drink}}\\ &\quad+0.736\,\mathbf h_{\text{tea}}\\ &=(0.100,\ 0.164,\ 0.736)\end{aligned}$$

Almost three quarters of the attention goes to the note of "tea". The machine is about to write "chai", and it is looking at "tea". Nobody told it to — the scores did.

Do this for every Hindi word and stack the weights as rows. You get the alignment map: one row for each word written, one column for each word read. In "I drink tea" → "मैं चाय पीता हूँ" (main chai peeta hoon) Hindi puts the verb at the end. So the bright cells do not run straight down the diagonal. They bend: चाय looks at "tea", and then पीता (peeta, "drink") looks back at "drink". The last word, हूँ (hoon, "am"), has no single English partner, so it spreads its attention over "I" and "drink".

Three ways to score a note. The dot product is not the only way to ask "how well does this note match?". Here are the three classic choices side by side.

An additive score by hand. Take $W_1=W_2=I$, $\mathbf v=(1,1)$, a decoder state $\mathbf q=(0.5,\ -0.5)$ and three notes $(1,0)$, $(0,1)$, $(1,1)$. Both matrices are $I$, so $W_1\mathbf h+W_2\mathbf q$ is just $\mathbf h+\mathbf q$: add the state to each note: $(1.5,\ -0.5)$, $(0.5,\ 0.5)$, $(1.5,\ 0.5)$. Squash each entry with tanh: $(0.905,\ -0.462)$, $(0.462,\ 0.462)$, $(0.905,\ 0.462)$. Now $\mathbf v=(1,1)$ simply adds the two entries: the scores are $(0.443,\ 0.924,\ 1.367)$, and softmax turns them into the weights $(0.195,\ 0.315,\ 0.490)$.

Rule of thumb. The dot product is the cheapest, and it is the one the transformer uses (with the $\sqrt d$ of §4). Additive scoring came first. Like general, it copes when the decoder state and the notes have different lengths: its matrices first bring both to one size.

The road ahead. The unit has four acts.

1. Look back instead of remembering (§1–§2): the decoder's look-back, and attention as a soft lookup.
2. Every word asks every word (§3–§7): self-attention by hand, why we divide by $\sqrt d$, one word with many meanings, an attention lab where you type your own sentence, and many heads.
3. Order and depth (§8–§9): three ways to tell the model where each word stands, and the transformer block.
4. How transformers learn (§10–§14): the mask that trains every next word at once; BERT, GPT and the encoder–decoder; the $n^2$ price; and a tiny transformer walked end to end.

The question. The decoder of §1 did not pick the best note. It took a blend of all of them. Why a blend — why not simply pick the winner?

First, give the three parts their names, with the help of a library. Your query is what you are looking for: "a cricket book for beginners". Every book has a label on its spine — its key. Inside the book is what you actually take home — its value. A strict librarian hands you the one book whose label matches best. A kind librarian photocopies a few pages from every book, more pages from the books whose labels match your wish better, and staples them into one bundle. That bundle is attention's answer.

In §1 each note played both parts at once: the decoder compared its state with the note (the note as a key) and then blended the notes themselves (the note as a value). Split the two parts and you have the general tool, the soft lookup.

A tiny example. The stall has two drinks. Chai costs ₹10, coffee costs ₹20. Your wish matches chai with a score of 2 and coffee with a score of 0. Turn the scores into shares with softmax (Unit 14):

$$\begin{aligned}\text{chai: }&\frac{e^{2}}{e^{2}+e^{0}}=\frac{7.389}{8.389}\approx0.881\\ \text{coffee: }&\frac{e^{0}}{e^{2}+e^{0}}=\frac{1}{8.389}\approx0.119\end{aligned}$$

The two shares add up to 1. Your mixed cup is 88% chai and 12% coffee, so its price is the same blend of the prices: $0.881\times10+0.119\times20\approx11.19$ — ₹11.19. A strict lookup would have said ₹10, full stop.

Where does a score come from? From the dot product of Unit 3: write the query and each key as lists of numbers, multiply them entry by entry and add up. Big when they point the same way, zero when they have nothing in common, negative when they disagree.

Strict against soft, side by side.

Rule of thumb. Inside a network, always go soft. If you want it nearly strict, make the scores big: a very picky soft lookup behaves like a dictionary, and it can still learn.

The question. In §1 the decoder looked back at the encoder's notes. But the encoder itself still read the sentence as a relay: the note written after "tea" knows about "I" only because "I" was passed along from hand to hand. Why keep the relay at all? What if every word could look at every other word directly?

That is self-attention: the words of one sentence attend to each other. In 2017 Vaswani and others built a whole translator from it, with no relay at all — the transformer, in a paper called "Attention Is All You Need". The relay of Unit 17 had two problems, and self-attention removes both.

1. The relay forgets. The first word's voice is handed along seven times to reach word 8. If each hand-over keeps about half of it, what arrives is $0.5^7\approx0.0078$ — less than 1%. With attention the path from word 1 to word 8 is one step long.
2. The relay is slow. Word 8 must wait for word 7, which waits for word 6. A computer with a thousand workers still has to do the steps one after another, so most of the workers sit idle. With attention all the questions are asked at once.

The price is the number of questions. With 8 words, each of the 8 asks all 8 (itself too): $8\times8=64$ questions. We come back to that price in §12.

Three views of every word. For every word to ask every other word, each word needs three things. Think of three badges it wears:

• a query $\mathbf q$ — "what am I looking for?"
• a key $\mathbf k$ — "what do I offer, so that others can find me?"
• a value $\mathbf v$ — "what do I hand over when someone picks me?"

All three are made from the same word. The word starts as one list of numbers $\mathbf x$ (its vector from the table of Unit 16), written as a row. Three learned matrices turn that one row into three:

$$\mathbf q=\mathbf x\,W_Q,\quad\mathbf k=\mathbf x\,W_K,\quad\mathbf v=\mathbf x\,W_V.$$

In this unit queries are orange, keys are blue, values are green and attention weights are gold. Training chooses the three matrices.

One word, worked by hand. Three words already have their keys and values, each a list of two numbers (so $d=2$):

$$\begin{aligned}&\mathbf k_1=(1,0),\ \ \mathbf k_2=(0,1),\ \ \mathbf k_3=(1,1)\\ &\mathbf v_1=(2,0),\ \ \mathbf v_2=(0,2),\ \ \mathbf v_3=(1,1)\end{aligned}$$

A word asks with the query $\mathbf q=(2,0)$. Four small moves:

1. Score each key with a dot product, then divide by $\sqrt d=\sqrt2\approx1.414$ (§4 says why): $\mathbf q\cdot\mathbf k_1=2$, $\mathbf q\cdot\mathbf k_2=0$, $\mathbf q\cdot\mathbf k_3=2$. Divided: $(1.414,\ 0,\ 1.414)$.
2. Turn scores into shares with softmax: $e^{1.414}\approx4.113$ and $e^{0}=1$, total $9.227$. Shares $\approx(0.446,\ 0.108,\ 0.446)$.
3. Blend the values with those shares: $0.446\,(2,0)+0.108\,(0,2)+0.446\,(1,1)$.
4. Read the answer: $\approx(1.337,\ 0.663)$.

The query points along the first direction, so it matched keys 1 and 3 equally well and key 2 badly. The answer leans towards $\mathbf v_1$ and $\mathbf v_3$.

All the words at once. In a sentence every word asks its own question. Write the shapes first (the habit of Unit 15). A sentence of $n$ words, each a list of $d$ numbers, is a table $X$ of shape $n\times d$. Stack the queries as the rows of $Q=XW_Q$, the keys as the rows of $K=XW_K$, the values as the rows of $V=XW_V$ — each $n\times d$. Then every lookup happens in one line:

$$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\Big(\frac{QK^{\mathsf T}}{\sqrt d}\Big)\,V.$$

Read it aloud, from the inside out. $QK^{\mathsf T}$: every question against every label — an $n\times n$ table of agreements. Divide by $\sqrt d$: calm the numbers down (§4). Softmax, row by row: each word shares out exactly 100% of its attention. Times $V$ ($n\times d$): each word collects its blend of gifts — $n$ new vectors, one per word. The table of shares is the attention matrix. With the three queries $(2,0)$, $(0,2)$ and $(1,1)$:

$$A\approx\begin{pmatrix}0.446&0.108&0.446\\ 0.108&0.446&0.446\\ 0.248&0.248&0.503\end{pmatrix}.$$

The answers are $(1.337,0.663)$, $(0.663,1.337)$ and $(1,1)$. The third query $(1,1)$ likes key 3 most, and its answer is exactly $\mathbf v_3$ — the other two values balance each other.

The relay against self-attention.

Rule of thumb. When the text fits, attention wins: nothing fades and nothing waits. The $n^2$ bill is what limits how long the text can be.

The question. Every score in §3 was divided by $\sqrt d$. What goes wrong without it — and why $\sqrt d$, not $d$?

A dot product adds up $d$ little products, one for each entry. More entries means a bigger pile. If the entries of $\mathbf q$ and $\mathbf k$ are random numbers of typical size 1, each little product is about size 1 too, and they add up like coin-toss steps: some push up, some push down, and they partly cancel. The pile's typical size — its spread, the standard deviation of Unit 14 — grows like $\sqrt d$:

$$\begin{aligned}d=2&:\ \sqrt2\approx1.41\\ d=64&:\ \sqrt{64}=8\\ d=512&:\ \sqrt{512}\approx22.6\end{aligned}$$

Why is that bad? Softmax only cares about gaps between scores. Take two scores 8 and 0, a normal gap when $d=64$. Softmax gives

$$\Big(\frac{e^8}{e^8+1},\ \frac{1}{e^8+1}\Big)\approx(0.99966,\ 0.00034).$$

That is nearly one-hot: the word listens to one other word and ignores the rest. Worse, the row is stuck. The slope of a share $p$ is $p(1-p)\approx0.00034$, almost zero, so training gets almost no blame through it (the fading blame of Unit 15).

Now divide by $\sqrt{64}=8$ first. The scores become 1 and 0, and softmax gives $(0.731,\ 0.269)$. The slope is $0.731\times0.269\approx0.197$ — nearly 600 times bigger. Dividing by $\sqrt d$ brings the spread back to about 1, whatever $d$ is.

Three choices for the same pair of scores, 8 and 0, with $d=64$.

Rule of thumb. Divide by $\sqrt d$ — the size of the spread. Too little and the rows freeze; too much and they go blind.

The question. In Unit 16 every word got one vector, the same in every sentence. So "bank" had to sit halfway between rivers and money, and "bat" halfway between cricket and caves. Can attention give a word a different vector in each sentence?

Yes — and you have already seen how. After self-attention a word's new vector is a blend of the values of the words it attends to. Change the neighbours and the blend changes. Watch it happen with numbers.

By hand. Take lists of just two numbers, made by hand: the first number says "how much about nature", the second "how much about money". So river $=(2,0)$, money $=(0,2)$, and bank $=(1,1)$ — right in the middle, because Unit 16 had to give it one vector for both meanings. To keep the arithmetic small, let every word use its own vector as its query, key and value ($Q=K=V=X$); a real model would first multiply by $W_Q$, $W_K$ and $W_V$.

Sentence one: "river bank". Bank asks with its query $(1,1)$:

$$\begin{aligned}\text{scores: }&(1,1)\cdot(2,0)=2,\quad(1,1)\cdot(1,1)=2\\ \div\sqrt2:\ &(1.414,\ 1.414)\ \Rightarrow\ \text{shares }(0.5,\ 0.5)\\ \text{bank}_{\text{new}}&=0.5\,(2,0)+0.5\,(1,1)=(1.5,\ 0.5)\end{aligned}$$

Sentence two: "money bank". The same steps give $0.5\,(0,2)+0.5\,(1,1)=(0.5,\ 1.5)$.

The same word now has two different vectors. The first leans towards nature, the second towards money. How different are they? Their cosine (Unit 3) is

$$\cos=\frac{(1.5)(0.5)+(0.5)(1.5)}{\sqrt{2.5}\,\sqrt{2.5}}=\frac{1.5}{2.5}=0.6,$$

an angle of about 53°. Before attention the two "bank"s were one and the same vector, with cosine 1. This is what contextual embeddings means: vectors computed from the whole sentence, not looked up once per word. ELMo (2018) built them with the two-way LSTMs of Unit 17; BERT (2018) built them with attention (§11).

One vector per word, against one vector per word in each sentence.

Rule of thumb. Static vectors are cheap and fine for "which words are alike?". Whenever the meaning depends on the sentence — which in real text is almost always — use contextual vectors.

The question. What does attention look like on a whole sentence — who looks at whom?

Give every word of the sentence its query, key and value, and work out every row of the attention matrix. With 9 words you get a $9\times9$ table. Row "it" says how much "it" listens to each word. Draw the table as a grid of squares, brighter where the share is bigger: that is an attention map.

The lab below does exactly the maths of §3, for any short sentence you type. Its word vectors are toy ones, made by hand — not taken from a trained model. Each word is a list of 12 numbers. Eight of them are labelled features you can read: thing, someone, action, "it"-word, "he/she"-word, describing word, place, small linking word. The other four say where the word stands (§8). The matrices $W_Q, W_K, W_V$ are hand-set too, so you can see why each share comes out as it does. But the arithmetic — dot products, $\sqrt d$, softmax, the blend — is exactly the real one.

In the lab, head 1 is set up so that an "it"-word asks "which word is a thing?", and every thing answers "I am a thing". So row "it" gives "ball" a share of about 0.766, and each of the other eight words about 0.029.

Look at the bars under the lab. Before attention, "it" is a pure "it"-word: its thing bar is 0. After attention it has collected about 0.77 of ball's value from head 1 (and a little more from head 2), so its thing bar has risen to about 0.82. The word "it" now carries part of the meaning of "ball" — the payoff of §5, in a real sentence. The same word in another sentence would pick up something else.

The question. A word often needs two answers at once: "which thing do I mean?" and "what came just before me?". One row of softmax gives one blend. How can one word ask two questions?

Run several attentions side by side. Each one is called a head. Each head has its own $W_Q, W_K, W_V$, builds its own attention map and makes its own blend.

Here is the trick that keeps it cheap. Do not give each head the full length $d$. Cut it into $h$ equal pieces. Each head's matrices make queries, keys and values of length

$$d_k=\frac{d}{h}.$$

Tiny numbers. A common size is $d=512$ with $h=8$ heads, so each head works with lists of length $512/8=64$. The eight answers, each of length 64, are glued side by side back into one list of length $8\times64=512$. A last matrix $W_O$ (size $512\times512$) mixes them.

The cost is the same as one big head. One big head needs three $512\times512$ matrices. Eight small heads need $8\times3$ matrices of size $512\times64$, which is exactly the same number of weights. Add $W_O$ and either way attention has $4d^2=4\times512^2=1\,048\,576$ weights.

Each head wears its own glasses. A head's matrices keep only some directions of the word vectors and throw the rest away — the projection of Unit 12. So the same words can be close together through one pair of glasses and far apart through another. Below, eight words have three hand-made features: is it a drink?, is it hot?, is it a sport? Three heads each keep just one feature ($d=3$, $h=3$, so $d_k=1$). Ask about "coffee". The drink head finds chai and lassi. The hot head finds the sun, coffee itself and chai. The sport head has nothing to ask, so it spreads its attention evenly.

One big head against many small ones ($d=512$).

Rule of thumb. Keep each head's size $d_k$ around 64: BERT-base uses $768/12=64$, and so does the small GPT-2.

The question. "Dog bites man" and "man bites dog" use the same three words. Can attention tell them apart?

Look back at the attention formula. It compares words by dot products and blends their values. Nowhere does it ask "which word came first?". So attention on its own is order-blind. Shuffle the words, and every word gets exactly the same answer as before — the answers just come out in the shuffled order.

A tiny check. Take toy vectors dog $=(2,0)$, man $=(0,2)$, bites $=(1,1)$, and let each word's query, key and value be its own vector. In "dog bites man" the word dog ends up as $(1.723,\ 0.277)$. In "man bites dog" it ends up as … $(1.723,\ 0.277)$. The same. Attention cannot tell who bit whom.

The fix: put the position into the vectors. Give the first word the tag $(1,0)$, the middle word $(0,0)$ and the last word $(0,1)$, and add the tag to each word's vector before attention. Now dog becomes $(2.967,\ 0.019)$ when it comes first and $(1.576,\ 1.284)$ when it comes last. Order finally matters.

Those three tags were made up for three words. Real models need tags for thousands of positions. There are three standard ways to make them.

1 · Learned position vectors. Keep a second lookup table, exactly like the word table of Unit 16, with one row per position: row 1 for the first word, row 2 for the second, and so on up to the longest text used in training (512 in BERT, 1 024 in GPT-2). Add the row to the word's vector. Training fills the table. It is simple — but the table has no row 513, so a model trained on 512 positions cannot read a longer text.

2 · Clock hands (sinusoidal tags). The first transformer built the tags from clock hands turning at different speeds. Picture a row of clocks. For position $p$, a clock of speed $\omega$ gives the two numbers $\sin(\omega p)$ and $\cos(\omega p)$ — where its hand points. The speeds are $\omega_i=1/10000^{2i/d}$. With $d=4$ there are two clocks, of speeds 1 and 0.01. Each tag lists the fast clock's two numbers, then the slow clock's two:

$${\small\begin{aligned}\text{PE}(1)&=(0.8415,\ 0.5403,\ \ 0.0100,\ 1.0000)\\ \text{PE}(2)&=(0.9093,-0.4161,\ \ 0.0200,\ 0.9998)\\ \text{PE}(3)&=(0.1411,-0.9900,\ \ 0.0300,\ 0.9996)\end{aligned}}$$

The fast clock tells near positions apart. The slow clock tells far positions apart — the hour hand and the minute hand together tell the time better than either alone. And here is the gem. Take the dot product of two neighbouring tags:

$$\begin{aligned}\text{PE}(1)\cdot\text{PE}(2)&\approx1.5403\\ \text{PE}(2)\cdot\text{PE}(3)&\approx1.5403\end{aligned}$$

The same number. For each clock, $\sin a\sin b+\cos a\cos b=\cos(a-b)$, and two positions $k$ apart always have their hands $\omega k$ apart. So $\text{PE}(p)\cdot\text{PE}(p+k)=\cos k+\cos(0.01k)$: it depends only on the distance $k$, never on where the pair sits.

3 · Turn the arrows (rotary positions, RoPE). Many modern models use the clock idea more directly. Instead of adding a tag, they turn the query and the key. Split each into pairs of numbers — little arrows. A word at position $m$ turns its query arrow by $m\theta$; a word at position $n$ turns its key arrow by $n\theta$. Turning is the rotation of Unit 4. The dot product of two arrows depends only on the angle between them, and turning both by the same extra amount changes nothing. So the score depends only on the gap $m-n$.

Tiny numbers. Take $\theta=30^\circ$ and $\mathbf q=\mathbf k=(1,0)$. Put the query at position 3 and the key at position 1: the arrows point at $90^\circ$ and $30^\circ$, the angle between them is $60^\circ$, and the score is $\cos60^\circ=0.5$. Move them to positions 5 and 3: arrows at $150^\circ$ and $90^\circ$, still $60^\circ$ apart, score 0.5. Positions 2 and 0: again 0.5.

The three ways, side by side.

Rule of thumb. For a new model, turn the arrows: nothing to learn, a turn for every position, and "two words back" means the same everywhere. (A model still reads best at the lengths it practised on; stretching it much further needs extra tricks.) Learned rows are the simplest choice when every text has a fixed maximum length.