#Attention

After embedding, every token is a vector, but each vector only knows about its own token. The vector for it is the same in every sentence, whether "it" means an animal, a street or the weather. To predict well, each token has to pull in information from the other tokens. Attention is the mechanism that does it, and it is the one idea that made transformers work.

#The intuition

Read this sentence:

The animal didn't cross the street because it was too tired.

What is tired? The animal. Now change one word:

The animal didn't cross the street because it was too wide.

Now "it" is the street. The words around "it" decide what it means.

So each token has to look at the others. But not at all of them equally: to make sense of tired, the word animal matters a lot and the hardly at all. A plain average of all earlier vectors would drown the useful word in noise. And which words matter depends on who is asking: tired needs a living thing, wide needs a place.

Attention solves this. Each token gives every earlier token (and itself) a weight, its share of attention, and the weights add up to 100%. The token's new vector is a blend of the vectors it looked at, mixed in those proportions.

Figure 5.1 · Where one attention head looks

Who is “it”?Previous wordSubject finderResting

…tired…wide

A coreference head: the last word looks back to find what “it” refers to. Hover, tap or use the arrow keys on a word on the left. Dots show information flowing from the attended words into it; greyed words are in its future and masked out.

Notice one honest detail. In a GPT-style model, it cannot see the word tired yet: that word comes later. So at the it position the head hedges between animal and street. The ambiguity is resolved at the last word, which can look back at everything. Information flows only forward in the text.

#Queries, keys and values

Where do the weights come from? The model has to decide, for every pair of tokens, how relevant one is to the other, using nothing but their vectors. It does this the way you search a library:

• You arrive with a query: "books about cats".
• Every book has a label on its spine, a key: "pets", "cooking", "Paris".
• You compare your query with each label and take home the contents of the books that match: their values.

Attention is a soft version of this search. Instead of taking only the best match, a token takes a little from every token it can see, and more from those whose key fits its query better. Every token plays all three roles at once:

• its query says what it is looking for (tired: "something that can get tired"),
• its key says what it offers to whoever is looking (animal: "I am a living thing"; street: "I am a place"),
• its value is the information it hands over when it is chosen.

Keys and values are separate because what makes a token easy to find is not the same as what is useful to pass on.

#A tiny example with numbers

Pretend the vectors have just two numbers: the first means "is a living thing", the second "is a place", like the character sheet in chapter 3. The query of tired is $q = (2, 0)$: it wants a living thing. It can see three keys, its own included:

The score is a dot product: multiply matching numbers and add them up, so animal gets $2 \times 1.5 + 0 \times 0 = 3$. It is large when the query and the key point the same way. A softmax (scores in, probabilities out) turns the scores $3, 0, 0$ into weights of about $91\%$, $4.5\%$ and $4.5\%$. The new vector for tired is

$$\begin{aligned} &0.91\, v_\text{animal} \\ +\ &0.045\, v_\text{street} \\ +\ &0.045\, v_\text{tired} \end{aligned}$$

almost pure information about the animal. Give the last word the query $q = (0, 2)$ instead, as wide would ("something that can be wide"), and the same keys hand $91\%$ to street. Same words, different query, different blend.

In a real model nobody decides what the numbers mean. Each token makes its three vectors by multiplying its own vector $x$ by three learned matrices,

$$q = x W_Q \qquad k = x W_K \qquad v = x W_V$$

and training adjusts $W_Q$, $W_K$ and $W_V$ until the matches they produce help predict the next token. The vectors have dozens or hundreds of numbers, and their directions rarely have names as clean as "living thing" (that part is simplified so you can do the sums by hand).

#All tokens at once

Every token runs this search at the same time, over the tokens of its own text, which is why it is called self-attention. Stack the queries into a matrix $Q$ (one row per token), the keys into $K$ and the values into $V$, and the whole computation is one formula:

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V$$

Read it from the inside out:

• $K^\top$ is $K$ transposed, flipped so its rows become columns, which makes $QK^\top$ a grid holding every query · key score at once: row $i$, column $j$ is how well token $i$'s query matches token $j$'s key.
• Dividing by $\sqrt{d_k}$, where $d_k$ is the length of each key vector, keeps the scores from getting too big (why, below).
• $M$ is the causal mask, the rule that a token may only look at itself and earlier tokens: $0$ where token $j$ is at or before token $i$, and $-\infty$ for the future. Since $e^{-\infty} = 0$, future tokens get exactly zero weight.
• Softmax turns each row of scores into weights that add up to 1, and multiplying by $V$ blends the values with those weights.

Why the mask? The model writes text left to right, so when it produces a word, the words after it do not exist yet. In training the whole text is processed at once, and the mask stops each position from peeking at the very words it is supposed to predict.

Now the real thing, step by step: four tokens with four numbers each, one head, and no position rotation to keep it small. Press play or step through with the arrows. Every number on screen is computed live from the ones before it, so at the scores step you can check the highlighted cell by hand.

Figure 5.2 · Self-attention, one step at a time

X · token vectors

d₀d₁d₂d₃

The0.2−0.40.10.9

·cat1.00.3−0.60.2

·sat−0.30.80.5−0.1

·down0.40.60.9−0.5

1/12

Four tokens, each already an embedding vector of 4 numbers (real models use thousands). Stacked, they form the matrix X.

Three things are worth noticing:

1. Everything is matrix multiplication. All queries are scored against all keys in one product. That is why attention runs so well on GPUs (chips with thousands of small cores that multiply in parallel), and why training handles every position of a sequence in parallel.
2. The output has the same shape as the input. Four vectors in, four vectors out, so attention layers can be stacked.
3. The weights depend on the input. A plain matrix multiplication, like the one by $W_Q$, applies the same numbers to every input. Here the mixing weights $A$ are computed from the tokens themselves. That is what lets "it" mean different things in different sentences.

#Why divide by √d?

Add up $d$ products of random numbers around $\pm 1$ and the total is typically about $\sqrt{d}$ in size: a dot product of two 4-number vectors lands around 2, of two 1,024-number vectors around 32. With big vectors, raw scores get large, and softmax of large numbers is almost one-hot: one key takes nearly 100%, the gradients for every other key (the signals training uses to adjust them) nearly vanish, and training stalls.

Figure 5.3 · Softmax saturation

Vector size d_k64

Scores range over 27.4 (a dot product of d random numbers spreads like √d = 8).

softmax(q·k)

key 1100.0%

key 20.0%

key 30.0%

key 40.0%

key 50.0%

key 60.0%

softmax(q·k ÷ √d)

key 149.5%

key 21.6%

key 36.6%

key 416.3%

key 510.2%

key 615.7%

Dividing by $\sqrt{d_k}$ keeps the scores around size 1 no matter how wide the head is.

#Many heads at once

One attention pattern can only express one kind of relationship. A token gets one set of weights that adds up to 100%, so if tired needs both the noun it describes and the word right before it, a single head has to split its attention and blurs both. So transformers run several heads side by side, each with its own $W_Q$, $W_K$, $W_V$. With a model width of $d_\text{model} = 4096$ and 32 heads, each head works in its own 128-dimensional slice.

Figure 5.4 · Four heads, one sentence

Who is “it”?

The

animal

didn’t

cross

the

street

because

it

was

too

tired

Previous word

The

animal

didn’t

cross

the

street

because

it

was

too

tired

Subject finder

The

animal

didn’t

cross

the

street

because

it

was

too

tired

Resting

The

animal

didn’t

cross

the

street

because

it

was

too

tired

Rows are queries, columns are keys (same word order). Hover a cell.

The outputs of all heads are concatenated back into one vector of width $d_\text{model}$ and multiplied by one more learned matrix, $W_O$, which mixes what the heads found:

$$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\, W_O$$

Researchers who look inside trained models do find heads with recognizable jobs: heads that attend to the previous token; induction heads that find an earlier copy of the current token and predict whatever followed it last time (the basis of copying patterns from the prompt); heads that link verbs to their subjects; and many heads that "rest" on the first token. Most heads, though, do things that are not easy to name.

Heads do not have to keep separate keys and values. In multi-query attention (MQA) all query heads share one key head and one value head; in grouped-query attention (GQA), used by most current open models, small groups of query heads share one. The queries stay diverse while the memory for keys and values shrinks several times over. Chapter 10 shows why that memory matters so much.

#One full pass, start to finish

Here is everything one attention layer does to a text of $n$ tokens, with the sizes of Llama 2 7B ($d_\text{model} = 4096$, 32 heads of 128 numbers each):

1. In: $n$ vectors of 4,096 numbers, one per token.
2. Project: each head multiplies them by its own $W_Q$, $W_K$, $W_V$, so every token gets a query, a key and a value of 128 numbers. With RoPE, queries and keys are rotated by their position here.
3. Score: each head compares every query with every key: an $n \times n$ grid of dot products, divided by $\sqrt{128} \approx 11.3$.
4. Mask and softmax: the future is set to $-\infty$, and softmax turns each row into weights that add up to 1.
5. Blend: each token's output in that head is the weighted average of the values it can see, 128 numbers.
6. Combine: the 32 head outputs are glued together ($32 \times 128 = 4096$) and multiplied by $W_O$.
7. Out: $n$ vectors of 4,096 numbers again, the same shape as the input. Each is an update for its token, carrying what that token gathered from the others; the next chapter shows how it is added back in.

Steps 2 to 5 are exactly what Figure 5.2 showed, only with one head and four numbers.

#The cost of attention

Every token is scored against every earlier token, so the work grows with the square of the sequence length $n$: twice the text, four times the scores. A 128,000-token context means about $8 \times 10^9$ scores per head per layer. Chapter 12 is about what this does to long contexts.

One family of fixes is simply to compute fewer scores. Instead of full causal attention, a head can use a sparse pattern:

Figure 5.5 · Which scores get computed

Full causalSliding windowSink + windowStrided sparse

window w4

70scores computed of 136 with full causal attention528.6Mat 128k tokens (w = 4096) of 8.6B with full causal attention

A window plus the very first token. Keeping that “attention sink” is what lets a windowed model stay stable on endless streams.

Hover a row to see what that token can attend to.

• Sliding-window attention lets each token see only the last $w$ tokens (Mistral 7B used $w = 4096$). Information from further back still arrives indirectly: each layer extends the reach by another $w$.
• Attention sinks. Remember the "resting" head that parks its weight on the first token? Softmax has to put its probability somewhere, and models learn to use the first token as a dump. If a windowed cache throws that token away, quality collapses. Keeping it (StreamingLLM) fixes this; chapter 11 shows how caches use the trick.
• Hybrid layouts. Several current models alternate: most layers use a local window, every few layers use full global attention. That keeps most of the quality at a fraction of the cost.

These patterns change what is computed. Other tricks keep full attention but compute it faster or store less: FlashAttention, the KV cache and grouped queries (chapter 10).

#Summary

• Attention lets each token gather information from earlier tokens, with weights that depend on the content.
• Each token makes a query (what it looks for), a key (what it offers) and a value (what it hands over), each by multiplying its vector by a learned matrix.
• Queries match against keys with dot products; softmax turns matches into weights; the weights average the values.
• Scores are divided by $\sqrt{d_k}$ to keep softmax from saturating, and a causal mask hides the future.
• Multi-head attention runs many such lookups in parallel and mixes them with $W_O$; GQA and MQA let heads share keys and values.
• Full attention costs $O(n^2)$, work that grows with the square of the length; sliding windows, attention sinks and sparse patterns compute fewer scores.

Attention is only half of a transformer layer. Next: the block that wraps it, and why it can be stacked a hundred times.