Unit 18 · Attention and Transformers — The Math Behind the Machine

The bottleneck, and the fix: look back

Imagine this

A minister gives a speech in English, and an interpreter must say it in Hindi. The first interpreter listens to the whole speech, holds it in her head, and then speaks. For a short speech this works. For a long one, the details are gone before she opens her mouth.

The second interpreter keeps notes, one line for every phrase. Before each Hindi word she glances back at her notes, and she looks hardest at the line that matters right now. She never has to remember the whole speech. She only has to know where to look.

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.

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.
Turn the scores into shares with softmax (Unit 14). The shares are positive and add up to 1. They are called the attention weights.
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".

Try: Press ▶ translate and watch the four Hindi words appear. Before each word the beams reach back to the notes, and one more row of the map lights up — see it bend at चाय and पीता. Click चाय and push how sure to 2: the share of "tea" climbs to 0.936. Then switch to no attention (Unit 17): every word gets the same summary, and there is no map at all.

word being written

how sure1

Why does this work?

Two reasons. First, nothing is squeezed. A long sentence simply has more notes, and every note stays exactly as it was written. Second, the blame has a short road. When the decoder writes a wrong word, backprop (Unit 15) sends the blame back through the context vector straight into the notes it used — one hop, in proportion to their shares. In the relay, the blame for an early word had to crawl back through every step, and it faded on the way (Unit 17).

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.

score	formula	extra weights	in plain words
dot (Luong, 2015)	$\mathbf q\cdot\mathbf h$	none	"how much do we agree?" — needs $\mathbf q$ and $\mathbf h$ of the same length
general (Luong, 2015)	$\mathbf q^{\mathsf T}W\mathbf h$	one matrix $W$	first re-express the note with $W$ , then agree
additive (Bahdanau, 2015)	$\mathbf v^{\mathsf T}\tanh(W_1\mathbf h+W_2\mathbf q)$	$W_1,\ W_2,\ \mathbf v$	a tiny one-layer network judges the match

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.

Trap

There is not one context vector for the sentence. There is a new one for every word the decoder writes: four Hindi words, four context vectors, four rows of the map. The notes stay the same; what changes is where the decoder looks.

The realization

$\begin{gathered}\alpha_{tj}=\operatorname{softmax}_j\big(\operatorname{score}(\mathbf q_t,\mathbf h_j)\big)\\ \mathbf c_t=\sum_j\alpha_{tj}\,\mathbf h_j\end{gathered}$

Before writing word $t$ , score every note $\mathbf h_j$ against the decoder's state $\mathbf q_t$ , turn the scores into shares $\alpha_{tj}$ , and blend the notes into the context $\mathbf c_t$ . Keep every note and look back, instead of carrying one summary and hoping it holds everything.

Pause & predict

In the worked example you make the decoder's state for चाय twice as long: $\mathbf q=(0,\ 1,\ 4)$ . What happens to the three attention weights?

Pause & predict

A 30-word English sentence is translated into 35 Hindi words, with attention. How many context vectors does the decoder build, and how many numbers are in the alignment map?

If you want the algebra · 1 proof, step by step

Prove it · the blame reaches every note in one hop

Claim. With dot-product scores $s_j=\mathbf q\cdot\mathbf h_j$ , shares $\alpha_j$ and context $\mathbf c=\sum_j\alpha_j\mathbf h_j$ , the slope of the context with respect to one note is $\partial\mathbf c/\partial\mathbf h_j=\alpha_j\big(I+(\mathbf h_j-\mathbf c)\,\mathbf q^{\mathsf T}\big)$ . The blame reaches note $j$ directly, scaled by its own share — however long ago the note was written.

Note

j

moves

\mathbf c

in two ways: directly, through its own term

\alpha_j\mathbf h_j

, and through the scores —

s_j

changes, so every share changes:

\frac{\partial\mathbf c}{\partial\mathbf h_j}=\alpha_jI+\sum_i\mathbf h_i\,\frac{\partial\alpha_i}{\partial\mathbf h_j}.

Only the score

s_j

contains

\mathbf h_j

, and

\partial s_j/\partial\mathbf h_j=\mathbf q^{\mathsf T}

The slope of softmax is

\partial\alpha_i/\partial s_j=\alpha_i(\delta_{ij}-\alpha_j)

(the drawer of §4). So

\begin{aligned}&\textstyle\sum_i\mathbf h_i\,\alpha_i(\delta_{ij}-\alpha_j)\,\mathbf q^{\mathsf T}\\ &=\big(\alpha_j\mathbf h_j-\alpha_j\mathbf c\big)\mathbf q^{\mathsf T}\\ &=\alpha_j(\mathbf h_j-\mathbf c)\,\mathbf q^{\mathsf T}.\end{aligned}

\sum_i\alpha_i\mathbf h_i

is the context

\mathbf c

itself.

Add the two parts:

\alpha_jI+\alpha_j(\mathbf h_j-\mathbf c)\mathbf q^{\mathsf T}=\alpha_j\big(I+(\mathbf h_j-\mathbf c)\mathbf q^{\mathsf T}\big)

. No long product of step-by-step slopes appears: one hop. ∎ In Unit 17 the blame for an early note passed through one slope matrix for every later step of the relay, and a long product of them fades (Unit 17). Here the note of "tea", with share 0.736, gets its blame at 0.736 strength, wherever it sits.

The road ahead. The unit has four acts.

Look back instead of remembering (§1–§2): the decoder's look-back, and attention as a soft lookup.
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.
Order and depth (§8–§9): three ways to tell the model where each word stands, and the transformer block.
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.

In one sentence: Instead of squeezing the sentence into one summary, the decoder keeps every note and, before each word it writes, scores the notes, turns the scores into shares and blends them — an interpreter glancing back at the speaker's notes, looking hardest where it matters now.

Free preview · Unit 18 of 20

That was section 1. The rest of the unit opens when you unlock it.

13 more sections and the practice arena — 18 widgets, 31 checks and 16 solved problems in the whole unit (this preview had 1 widget and 2 checks).

Unlock Unit 18

2
Attention is a soft lookup
3
Self-attention: every word asks every word
4
Why we divide by the square root of d
5
The payoff: one word, many meanings
6
The attention lab: who looks at whom
7
Many heads: many questions at once
8
Where is each word? Three ways to add position
9
The transformer block
10
No peeking: the mask, and every guess at once
11
Three families: BERT, GPT and the encoder–decoder
12
The price: every pair of words
13
A tiny transformer, whole
14
What to carry forward
15
Practice arena — sixteen problems, solved in full

Unlock this unit for ₹299, or all seven paid units for ₹999 — one-time payment, full refund within 7 days. See pricing.

The bottleneck, and the fix: look back

That was section 1. The rest of the unit opens when you unlock it.

Attention is a soft lookup

Self-attention: every word asks every word

Why we divide by the square root of d

The payoff: one word, many meanings

The attention lab: who looks at whom

Many heads: many questions at once

Where is each word? Three ways to add position

The transformer block

No peeking: the mask, and every guess at once

Three families: BERT, GPT and the encoder–decoder

The price: every pair of words

A tiny transformer, whole

What to carry forward

Practice arena — sixteen problems, solved in full