Understanding Attention

Imagine reading the sentence "The animal didn't cross the street because it was too tired." To understand "it", you look back over the sentence and decide that "animal" is the relevant word, not "street". If the sentence ended "...because it was too wide", you would pick "street" instead. Understanding a word means gathering information from the right other words, and which words are right depends on the context.

Attention gives every token a way to do exactly that. For each token, it works out how relevant every other token is, then builds a new vector for the token from a relevance-weighted mix of the others. The token's new vector now carries context.

Figure 6.3.1. Bidirectional illustration with invented attention weights. Later tired or wide can affect an earlier position only when the mask permits it. A causal model at it cannot see those future words. Attention mixes projected values and is not a guaranteed explanation of a decision.

The rest of this lesson explains how a network computes "how relevant", and how it uses the answer.

First, let's see precisely why we need a mechanism like this. Most layers you know, such as linear layers, activation functions and the feed-forward network inside a Transformer block, process each token's vector on its own. Whatever the surrounding tokens are, the same input vector produces the same output vector.

Here is the word "bank" in two sentences, "the river bank" and "money in the bank", passed through such a per-token layer:

The two outputs for "bank" are identical (True). However many per-token layers we stack, "bank" by the river and "bank" with the money will always be processed identically, because no information ever flows between positions. This is the static-representation problem from Lesson 5.3, seen from inside a network.

To build context, tokens must be able to read from each other. That is the one job of attention.

Here is the whole mechanism for a sequence of T tokens, each represented by a vector. We call it self-attention because the tokens attend to other tokens of the same sequence.

1. Project. From each token's vector, compute three new vectors with three different learned linear layers: a query q, a key k and a value v.
2. Score. For every pair of tokens (i, j), compute a relevance score: the dot product of token i's query with token j's key.
3. Scale. Divide every score by √d, where d is the length of the query and key vectors.
4. Mask (decoders only). Block the scores for pairs that are not allowed, such as future tokens.
5. Normalise. Apply softmax across each row, so that each token's weights over all the tokens are positive and add up to 1.
6. Mix. Each token's output is the weighted average of all the value vectors, using its weights.

Written compactly, with the queries, keys and values for all tokens stacked as rows of matrices Q, K and V:

Don't worry if the formula looks dense: Sections 4 to 8 take it apart one step at a time, with numbers you can check by hand.

Figure 6.3.2. Project Q, K and V; compare queries with keys; scale by √dₖ; apply the mask before row-wise softmax; then mix values. Each output is a weighted sum of projected value vectors.

Why three different vectors per token? Think of a library search.

• The query is what you are looking for: the search you type.
• Each book has a key: the label or catalogue entry that describes it, used only for matching.
• Each book also has its value: the actual content you take away once you decide it is relevant.

Matching uses queries against keys; what you collect is values. Separating them lets a token look for one thing while offering something different. For example, a verb might look for its subject (its query), describe itself as an action (its key), and pass on information about tense (its value).

In a Transformer, all three are produced by learned linear layers from the token's vector, so the network learns what to look for, how to describe itself and what to pass on. To keep the arithmetic visible, we will use a hand-made example with 4 tokens, "the cat sat down", and query, key and value vectors of just 2 numbers each. The numbers are chosen so that "sat" and "down" look for something that "cat" offers, as a verb might look for its subject.

The score between token i and token j is the dot product of i's query with j's key: multiply matching entries and add. A large score means "j's key matches what i is looking for". We compute all 16 scores with two loops, then check against the single matrix product Q Kᵀ.

Read the score grid row by row; each row is one token asking "who is relevant to me?".

• The row for "sat" is [0, 4, 1, 0]. Its query [2, 0] matches the key of "cat", [2, 0], strongly: 2 × 2 + 0 × 0 = 4. The other keys match much less.
• The row for "down" is similar, [0.5, 4, 1.25, 0.25]: it is also most interested in "cat".
• "the" and "cat" share the query [0, 1], so their rows are identical: [1, 0, 0.5, 0.5].

The loops and Q @ K.T give exactly the same grid (True). In practice we always use the matrix product: one operation computes every score at once, which is exactly the parallelism that makes Transformers fast.

Scores can be any size, but we want weights: positive numbers that add up to 1 for each token, so that each token divides 100% of its attention among the others. Softmax does exactly that. But before softmax, the scores are divided by √d. Why?

A dot product adds up d products. If the entries of queries and keys are random numbers with a typical size of about 1, then the more numbers you add, the larger the typical result: its spread grows like √d. Let's measure that with 10,000 random query–key pairs.

The spread (standard deviation) of the raw dot products is 1.99 for d = 4, 8.01 for d = 64 and 22.45 for d = 512, growing like √d (√4 = 2, √64 = 8, √512 ≈ 22.6). After dividing by √d, it is 1.00, 1.00 and 0.99: the same at every size.

Why does the size of the scores matter? Because softmax exaggerates differences. If scores are large, the largest one swallows almost all the weight. Here are 8 scores with d = 512, passed through softmax with and without scaling:

Without scaling, softmax puts a weight of 1.000 on a single key and 0.0 on the other seven: the entropy is 0.000, which means no uncertainty and no blending at all. With scaling, the largest weight is 0.442 and the rest is shared among several keys, with an entropy of 1.480 (the maximum possible for 8 keys is 2.079).

This matters for learning, not just for the output. When softmax is that extreme, small changes to the scores barely change the weights, so the gradients flowing back through the softmax are close to zero, and the queries and keys stop learning. Dividing by √d keeps scores at a sensible size however wide the vectors are, so attention can stay soft and trainable.

Now we apply both steps to the "the cat sat down" scores. Here d = 2, so we divide by √2 ≈ 1.414.

Every row now adds up to 1. "Sat" gives 0.808 of its attention to "cat" and spreads the rest thinly. "Down" gives 0.771 to "cat". "The" and "cat", with their weak, similar scores, spread their attention more evenly: their largest weight is only 0.345.

Common mistake: applying softmax over the wrong dimension. The weights for one query must sum to 1 across the keys, so softmax goes over the last dimension (dim=-1) of the score matrix. If the weights for one token add up to something other than 1, check the dimension.

The final step uses the weights to mix the values. Each token's output is the weighted average of every token's value vector, using that token's row of weights. In matrix form, it is simply weights @ V.

The output for "sat" is [0.168, 0.928]. Our hand calculation, 0.048 × the value of "the" + 0.808 × the value of "cat" + 0.097 × the value of "sat" + 0.048 × the value of "down", gives the same vector. Since the value of "cat" is [0, 1], the output for "sat" has become mostly "cat": the verb has pulled in information about its subject. "The" and "cat", whose weights are more evenly spread, get a more blended output, [0.709, 0.534].

This is how attention creates contextual representations. Each output vector is no longer a description of a token on its own, but of the token together with what it found relevant. In a real Transformer, this output is added back into the token's residual stream (Lesson 6.1), so the token carries this context into the next layer.

Visual walkthrough: Mix values with causal attention.

Choose each query position in turn. Masked future tokens receive zero weight; allowed weights sum to one.

These supplied numbers illustrate the calculation. The web app adds a paused, step-by-step animation of the same example.

So far, every token could attend to every other token, including the ones after it. That is bidirectional attention, used in encoders. A decoder that generates text must not let a token see the future, because during generation the future has not been written yet. If it could peek during training, it would learn to "predict" the next token by simply reading it, and would fail completely when generating.

The fix is the causal mask from Lesson 6.1. Before softmax, we replace every score where the key comes after the query with −∞ (minus infinity). Since e to the power of −∞ is 0, softmax gives those positions exactly zero weight, and the remaining weights still add up to 1.

The masked score grid has −∞ above the diagonal, and the weights there are exactly 0. Each row still sums to 1. Notice how the rows change:

• "the" (position 0) can only see itself, so it gives itself all of its attention: 1.000.
• "cat" (position 1) can see "the" and itself: 0.670 and 0.330.
• "sat" (position 2) still gives most attention to "cat" (0.848), because "cat" comes before it.
• "down" (position 3) can see everything, so its row is unchanged from the unmasked version.

Key terms

• Self-attention: attention in which a sequence's tokens attend to tokens of the same sequence.
• Query, key, value: three learned projections of each token: what it looks for, how it is matched, and what it passes on.
• Attention score: the dot product of a query with a key, divided by √d.
• Attention weights: the softmax of the scores; each token's weights sum to 1.
• Causal mask: setting scores for future positions to −∞ so their weights become 0.

A single attention pattern can only express one kind of relevance at a time. But a token often needs several kinds of information at once: the previous word, the subject of the sentence, an earlier mention of the same name, the matching opening bracket. Multi-head attention runs several attention operations, called heads, side by side, each with its own query, key and value projections, so that each can learn a different kind of relationship.

To keep the cost the same, the model's vector is split between the heads rather than copied. With d_model = 64 and 4 heads, each head works with 16 numbers:

A (1, 10, 64) tensor is viewed as (1, 4, 10, 16): 4 heads, each seeing all 10 tokens with 16 numbers per token. Each head computes its own (10, 10) grid of weights. Afterwards, the heads' outputs are put back side by side into (1, 10, 64), and a final linear layer combines them. Reshaping back and forth loses nothing (True).

To see why several heads help, let's build two heads by hand, each designed to capture a different relationship, for the words "to be or not to be". Each token's input vector contains two parts: which word it is, and which position it is at. We set the query and key weights directly:

• Head A turns each token's position into a query for the previous position.
• Head B turns each token's word into a query for the same word.

The rows confirm the designs. Head A's row for "not" (position 3) puts 0.98 of its weight on position 2, the previous word "or". Head B's row for the second "be" (position 5) splits its weight between itself and the earlier "be" at position 1, 0.493 each: it found the previous copy of its own word. Here are both heads as heatmaps:

Head A forms a clear diagonal one step below the main diagonal: every token looks at its predecessor (the first token, having no predecessor, looks at itself). Head B lights up only where a word matches an earlier copy of itself. The two heads read the same sentence and gather completely different information, and the model receives both.

In a real Transformer, nobody designs the heads. Their weights are learned, and trained models do develop specialised heads like these. Patterns very much like Head A, "attend to the previous token", are among the most commonly found in trained language models. We will look at the heads of our own trained model in Lesson 6.6.

Key terms

• Head: one attention operation with its own query, key and value projections.
• Multi-head attention: several heads running in parallel on slices of the model vector, combined by a final linear layer.
• d_head: the vector length per head, d_model / n_heads.

Look back at the formula in Section 3. Nothing in it refers to where a token sits. Scores depend only on the vectors themselves. So if you shuffle the input tokens, every token still finds exactly the same partners with exactly the same weights, and the outputs are simply shuffled in the same way. Let's verify that with PyTorch's own attention layer:

Without position information, shuffling the 5 input tokens gives exactly the same outputs, in the shuffled order (True). To attention alone, "dog bites man" and "man bites dog" would be the same bag of words, as in Lesson 5.4. But when we add a vector for each position before attention (and the positions stay put while the tokens move), shuffling the tokens changes the outputs (False). The position vectors have broken the symmetry: the model can now tell orderings apart.

Transformers generally supply position through learned vectors, sinusoidal encodings, or mechanisms inside attention such as RoPE. The causal mask also carries some ordering information, so the unmasked permutation argument does not apply unchanged to causal attention. There are two common ways:

• Learned position embeddings: a table with one learnable vector per position, as in Lesson 6.2. This is what our model uses.
• Sinusoidal positional encodings: fixed vectors made from sine and cosine waves of many different frequencies, as in the original Transformer. Each dimension oscillates at its own speed, so every position gets a unique pattern.

Let's compute the sinusoidal version for 64 positions and look at it.

In the left panel, each column is one position's 64-number vector. The top dimensions oscillate quickly from position to position, and the lower ones slowly, rather like the hands of a clock: the second hand tells you small differences, the hour hand large ones. Together they give every position a distinct pattern.

The right panel shows a useful property: nearby positions get similar vectors. The cosine similarity of position 10 with itself is 1.0, with position 11 it is 0.966, with position 15 it is 0.734, and it tends to fall further with distance (0.59 at position 30, 0.49 at position 60), although not perfectly smoothly. So the encoding carries a sense of "near" and "far", which attention can use.

Common mistake: forgetting positional information when building a Transformer from scratch. The model will still train, but unmasked attention is permutation-equivariant, while causal attention has only the ordering information supplied by its mask, and it will plateau at a much worse loss than it should.

Attention's great strength, that every token compares itself with every other token, is also its main cost. With T tokens, each head computes a T × T grid of scores. Double the sequence length and that grid becomes four times larger.

With our context of 64 tokens, each head's score grid has 4,096 entries and needs 16 KB of memory at 32-bit precision. At 1,024 tokens it has 1,048,576 entries (4 MB): 256 times as many as at 64. At 8,192 tokens it needs 256 MB, per head, per layer, per sequence. At 131,072 tokens a single grid would need 65,536 MB, 64 GB, more memory than most GPUs have, and that is for just one head of one layer.

This quadratic cost is the main reason context windows were limited for many years, and a great deal of engineering goes into managing it. Two key ideas are worth knowing by name:

• Memory-efficient attention (for example FlashAttention) computes the same result in small tiles without ever storing the full T × T grid. PyTorch's F.scaled_dot_product_attention, which you will use in Section 12, can use such methods automatically.
• Restricted attention patterns let each token attend to only part of the sequence, such as a sliding window of recent tokens, trading some reach for a cost that grows more slowly.

Now we can put everything together into the attention module that our model will use in Lesson 6.5. It takes a batch of shape (B, T, d_model) and returns the same shape. Read the comments beside each line for the shapes.

• One linear layer produces queries, keys and values at once (3 * d_model outputs), which we split into three.
• view and transpose split the model vector into heads: (B, T, D) becomes (B, H, T, d_head).
• The core is the formula from Section 3: scores, scaling, causal mask, softmax, weighted sum of values. All heads are computed at once, because the head dimension is treated just like the batch dimension.
• Merge and project: the heads are put back side by side, and the output projection proj combines them.

Input and output are both (2, 12, 64), and the weights have shape (2, 4, 12, 12): for each of 2 sequences and 4 heads, a 12 × 12 grid. The module has 16,640 parameters: 64 × 192 + 192 for the combined query–key–value layer, plus 64 × 64 + 64 for the output projection.

Any implementation from scratch should be checked against a trusted one. PyTorch provides F.scaled_dot_product_attention, which computes the core of attention (in a fast, memory-efficient way) given queries, keys and values. We feed it the same queries, keys and values as our module and compare:

Our implementation matches PyTorch's (True), with a maximum difference of 1.2e-07, which is just floating-point rounding. And no weight sits above the diagonal (True), so the causal mask works. In real projects you would normally call F.scaled_dot_product_attention directly, because it is faster and uses less memory; writing it yourself once is how you understand what it does.

Finally, let's look at the weights our untrained module produces for the opening of Romeo's famous line, "ROMEO:\nBut, soft! what light", using random untrained embeddings.

The triangle shape is the causal mask at work: nothing above the diagonal. Within the triangle, the patterns are noisy and unfocused. For the last position, head 0's largest weight is 0.168, compared with 0.036 if the weight were spread evenly over all 28 positions. There is some structure, from the random weights, but nothing meaningful: an untrained module has not yet learned what to look for. In Lesson 6.6 we will draw the same kind of picture for the trained model and compare.

Layers that process each token on its own cannot build context: "bank" gave identical outputs in two different sentences. Self-attention fixes this in three steps: compare, normalise, average.

Each token produces a query (what it looks for), a key (how it is matched) and a value (what it passes on). Scores are query–key dot products, divided by √d, because otherwise their spread grows with the vector size (22.45 at d = 512) and softmax collapses onto one key (weight 1.000, entropy 0.000). Softmax turns each token's scores into weights that sum to 1, and each output is the weighted average of the values. In our example, "sat" put 0.808 of its attention on "cat" and came out mostly "cat". A causal mask sets future scores to −∞, so decoders never see ahead.

Multi-head attention splits the vector among several heads, each free to capture a different relationship, such as "previous word" or "earlier copy of me". Because attention alone ignores order, we add positional information, learned or sinusoidal. Attention's power has a price: each head's score grid grows with T², from 16 KB at 64 tokens to 64 GB at 131,072.

Our CausalSelfAttention module, with 16,640 parameters, matches PyTorch's scaled_dot_product_attention to within 1.2e-07. It is the attention layer we will use in our model.

Figure 6.3.3. Learn Q, K and V from the full 64-coordinate input, then reshape into four 16-coordinate heads. Concatenate their outputs and apply an output projection. Heads do not receive disjoint slices of the original unprojected input.