Attention
Lesson 3 of 8
3Attention
The big pictureAttention answers one question for every word: "which other words in the sentence matter to me now, and how much?". Each word asks (query), each word advertises what it has (key), and the answer is a blend of their values in proportions set by the match. This one idea is the heart of every modern language model.
What the picture shows: These are the weights of one row from the example: "sat" gives itself 51% because its key [1, 2] matches its query [1, 1] best. Notice they sum to 100%, so the output is just a weighted average of the values.
Map of this part
The concepts in order, click any one to jump to it:
Two ways into the same idea
When you read "the cup fell off the table because it was on the edge", you know at once that "it" means the cup. Attention gives the model that ability: each word asks the others "which of you matters to me?" and takes from each in proportion to its importance.
Each token has three vectors: query Q, key K and value V. The weights softmax(Q·Kᵀ / √d) set how much each token takes from the others' values, and the output is a weighted sum of values. A causal mask stops a token seeing what comes after it in generative models.
Before You Start
🧭Embeddings
You know every word is a vector, and that the dot product of two vectors measures how alike they are.
🎲Probability, practically
You know softmax: it turns any scores into positive shares that sum to 1.
The Core Idea
In this lesson you build attention step by step: query, key and value, then scores and blending, then why we divide by √d, then the mask, and finally its cost.
1Query, key and value
Each word is multiplied by three matrices, and three vectors come out. The query says what it looks for, the key says what it offers, and the value is what it hands over if chosen.
Look at the picture: the same word plays three roles. When it asks it uses its query, when it is asked its key is compared, and when it is chosen it gives its value.
2Scores, softmax, blend
We multiply every query by every key and get a table of scores, words by words. Then softmax turns each row into shares summing to 1, and we blend the values in those shares.
import numpy as np
words = ["the", "cat", "sat", "down"]
x = np.array([[1, 0, 0],
[0, 1, 0],
[0, 1, 1],
[0, 0, 1]], dtype=float) # one 3-number vector per word
Wq = np.array([[0, 1], [1, 0], [0, 1]], dtype=float)
Wk = np.array([[1, 0], [1, 1], [0, 1]], dtype=float)
Wv = np.array([[1, 0], [0, 1], [1, 1]], dtype=float)
Q, K, V = x @ Wq, x @ Wk, x @ Wv
scores = Q @ K.T / np.sqrt(2)
weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
out = weights @ V
np.set_printoptions(precision=2, suppress=True)
print("Q:\n", Q)
print("K:\n", K)
print("scores = Q K^T / sqrt(2):\n", scores)
print("weights (softmax of each row):\n", weights)
print("output = weights V:\n", out)
Read the third row, the word sat: its query is [1, 1], and the key of sat itself is [1, 2], so the product is 3, and after dividing by √2 it is 2.12, the highest in the row. After softmax it takes 51% of its attention, and the rest is spread out. Its output is a blend of the four values in those proportions. Compute the first row by hand and check it matches.
3Why divide by √d
In the table above we divided by √2 because the width is 2. This small division becomes very important as the width grows.
With a width of 128 or more, dot products become large numbers, so softmax gives almost all the weight to one word and the gradients for the rest vanish. Dividing by the square root of the width brings the numbers back to a sensible range.
4The causal mask
A model that generates text must not see the words after it. So we put −∞ above the diagonal before softmax, and the future's weight becomes zero.
We make all scores equal so we see the mask's effect alone.
import numpy as np np.set_printoptions(precision=2, suppress=True) T = 4 scores = np.zeros((T, T)) # equal scores, to see the mask alone mask = np.triu(np.ones((T, T), dtype=bool), k=1) # True above the diagonal = the future scores[mask] = -np.inf weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True) print(weights)
The first word sees only itself, the second splits its attention between itself and the first, and so on. Without this mask the model learns during training to "read" the next word from the input instead of predicting it, so the loss drops amazingly and then it fails completely at generation time.
Look at the picture: the first word sees only itself, and the fourth sees all four. Each word sees itself and what came before, nothing more.
5The quadratic cost
Every token is compared with every token, so doubling the context length means four times the compute and memory in attention.
for n in [512, 2048, 8192, 32768, 131072]:
pairs = n * n
mb = pairs * 2 / 1e6 # one fp16 score per pair, one head, one layer
print(f"{n:>7} tokens {pairs:>17,} pairs {mb:>10,.0f} MB per head per layer")
From 8 thousand to 32 thousand tokens, four times the length, memory goes from 134 to 2,147 megabytes, sixteen times, for one head in one layer. Modern models do not store this matrix in full thanks to algorithms like FlashAttention that compute it in tiles, but the compute itself stays quadratic. That is why long context is slow and costly even where it is available.
How It Works, Step by Step: how attention is computed for one sentence
Multiply each word's vector by three matrices Wq, Wk and Wv to get its query, key and value.
Multiply every query by every key and divide by √d, giving a table of scores of size words by words.
In generating models you put −∞ above the diagonal, then apply softmax to each row so it sums to 1.
Multiply the weights by the values, so each word's output is a blend of the words' values in its attention shares.
Worked example
We will compute attention for a four-word sentence with small whole numbers you can check on paper. Each word is a vector of three numbers, and three weight matrices turn it into a query, a key and a value of width 2.
Example: The same in PyTorch
In real code you do not write the softmax yourself. PyTorch has one function that picks the fastest implementation for your hardware.
import torch import torch.nn.functional as F q = torch.randn(1, 4, 16) # (batch, tokens, width) k = torch.randn(1, 4, 16) v = torch.randn(1, 4, 16) out = F.scaled_dot_product_attention(q, k, v, is_causal=True) print(out.shape) # torch.Size([1, 4, 16])
is_causal=True applies the same mask we wrote. The same equation, with an implementation that uses FlashAttention when available.
Try it without code
Myth vs Reality
| Myth | Reality |
|---|---|
| An attention map explains why the model said what it said. | It is blend proportions in one layer and one head out of dozens. On its own it proves no cause. |
| Dividing by √d is a cosmetic detail. | Without it scores grow with the width, softmax gives almost all the weight to one word, and the gradients for the rest vanish. |
| A longer context costs only in proportion to its length. | Attention compares every token with every token, so doubling the length quadruples the work: from 8 thousand to 32 thousand tokens memory went from 134 to 2,147 MB for one head. |
Try it yourself
What is the sum of each row in the attention weight matrix?
Softmax turns each row into proportions summing to 1, so the output is a weighted average of the values. Even with the mask, the blocked words get zero and the rest sum to 1.
Doubling the context from 4 thousand to 8 thousand tokens multiplies attention compute by how much?
The number of pairs is n², so doubling n multiplies pairs by 4. The rest of the model (the linear layers) only doubles, which is why attention dominates cost at long context.
Quick review cards
softmax(Q Kᵀ / √d) V: every query against every key, turned into proportions summing to 1, then used to blend the values.−∞ above the diagonal before softmax, so the future's weight is zero.Check Your Understanding
In our example, why are the first and last rows identical?
Because the and down have the same query [0, 1] (see Q), so they match the keys with the same scores, get the same weights and the same output. In attention, two words with the same query see the sentence exactly the same way, which is one reason we need position encoding.
Why does "down" get only 0.14 in the second row, less than the other three?
The query of "cat" is [1, 0] and the key of "down" is [0, 1], so their product is zero, the lowest score in the row. The other keys all start with 1, so each scored 0.71 after dividing by √2.
Resources for this lesson
Grouped by type. Start with the official docs for precision, or an article or video for the explanation.
Official docs2
Articles3
Done with this lesson?
Sign in to save your progress.