Attention Explained With Real Numbers: Scaled Dot-Product and Multi-Head

Every modern language model is built around one operation, and you can compute it with a pencil. Today we take a tiny three word sentence and work out attention completely by hand, every score, every weight, every output. By the end, the famous formula will feel like arithmetic you already know.

• Every transformer runs on it
• We use real numbers

Consider the word bank in a sentence about a river, and again in a sentence about money. The embedding from lesson 54 is the same for both. Attention lets each token look at the other tokens and mix in information from the ones that matter, so the meaning becomes context dependent.

• Words depend on context
• Each token must look around

We will build queries, keys and values, derive the one over square root of d k scaling, add masks, split into multiple heads, track tensor shapes, and write the code in NumPy and PyTorch. We do not cover the full transformer block here. That is the next lesson.

• Scores, softmax, values
• Scaling, masks, heads
• Code and complexity

Think of a library. You walk in with a query, a description of what you want. Every book has a key, a label describing its contents, and a value, the content itself. A normal dictionary returns the single book whose key matches exactly. Attention does a soft version of this.

• Query: what I want
• Key: what I offer
• Value: what I give

Instead of returning one book, attention compares the query to every key, turns the similarity into weights that add up to one, and returns a blend of all the values. Close matches contribute a lot, poor matches contribute almost nothing. And because it is all smooth arithmetic, we can train it with gradient descent.

• Hard lookup returns one item
• Soft lookup blends all items

In self attention, every token plays all three roles. Each token embedding is multiplied by three different learned matrices, W Q, W K and W V, producing its own query, key and value. The matrices are the learned part. The lookup itself has no parameters.

• Each token makes all three
• Three learned matrices

Here is our example. Three tokens, each with a two dimensional embedding. Token one is one, zero. Token two is zero, one. Token three is one, one. Stack them as rows and we get the matrix X, with shape three by two.

• 3 tokens, 2 dimensions
• X is a 3 by 2 matrix

Let W Q have rows one, zero and one, one. Multiply each token row by it. Token one gives one, zero. Token two gives one, one. Token three is one, zero plus one, one, which is two, one. So Q has rows one zero, one one, and two one.

• W_Q rows: (1,0) and (1,1)
• Q = X W_Q

For the keys, let W K have rows one, one and zero, one. Token one gives one, one. Token two gives zero, one. Token three gives one, one plus zero, one, which is one, two. So K has rows one one, zero one, and one two.

• W_K rows: (1,1) and (0,1)
• K = X W_K

For the values, let W V be diagonal, with two and one on the diagonal. It doubles the first coordinate and keeps the second. The result is V with rows two zero, zero one, and two one. We now have all three matrices, each of shape three by two.

• W_V = diag(2, 1)
• V rows: (2,0), (0,1), (2,1)

Why not use the embedding directly for all three roles? Because the features that make a token a good match for a question are not the same as the features it should hand over once chosen. Separate projections let the model learn what to ask, what to advertise, and what to deliver.

• Roles differ
• Same token, three views

Now the scores. Entry i, j is the dot product of query i with key j, which, as in lesson 3, measures alignment. For query one, one zero, the scores are one, zero, one. For query two, one one, they are two, one, three. For query three, two one, they are three, one, four.

• Entry (i,j) = q_i dot k_j
• Shape 3 by 3

Next we divide every score by the square root of d k. Here d k is two, so we divide by one point four one four. Focus on token two. Its scores two, one, three become one point four one four, zero point seven zero seven, and two point one two one. We will justify this step shortly.

• d_k = 2, so divide by 1.414
• Row two: 1.414, 0.707, 2.121

Softmax turns scores into weights. First, exponentiate each score. E to the one point four one four is about four point one one. E to the zero point seven zero seven is about two point zero three. E to the two point one two one is about eight point three four. Exponentiating keeps everything positive and amplifies differences.

• e^1.414 = 4.113
• e^0.707 = 2.028
• e^2.121 = 8.339

The sum is about fourteen point four eight. Divide each by it. We get zero point two eight four, zero point one four zero, and zero point five seven six. They add to one. So token two pays about fifty eight percent of its attention to token three, and only fourteen percent to itself.

• Sum = 14.48
• Weights: 0.284, 0.140, 0.576

Now we blend the values. Zero point two eight four times two, zero gives zero point five six eight, zero. Zero point one four times zero, one gives zero, zero point one four. Zero point five seven six times two, one gives one point one five two, zero point five seven six. Add them: the output is one point seven two, zero point seven two.

• 0.284·(2,0) + 0.140·(0,1) + 0.576·(2,1)
• Output: (1.72, 0.72)

Doing this for all three tokens at once is just matrix multiplication. The weight matrix A, shape three by three with rows summing to one, multiplies V, shape three by two, giving an output of shape three by two. Same shape as the input, but each row is now informed by its context. That is the whole formula.

• Output = softmax(QK^T / sqrt(d_k)) V
• Row-wise softmax

Why divide by the square root of d k? It looks arbitrary, but it fixes a real problem. Real models use d k of sixty four or more, not two. As the dimension grows, dot products get bigger, and that quietly breaks softmax. Let us derive exactly how.

• Why not just QK^T?
• Dimension changes the scale

Assume each component of q and k is independent, with mean zero and variance one, roughly what good initialization gives. A single product q i times k i has mean zero. Its variance is the expectation of q i squared times k i squared, which is one times one, so one.

• q_i, k_i independent
• Mean 0, variance 1
• Var(q_i k_i) = 1

The dot product is a sum of d k such products. For independent terms, variances add. So the variance of the dot product is d k, and its standard deviation is the square root of d k. At d k equals sixty four, typical scores spread about eight units either side of zero. At five hundred twelve, about twenty three.

• Dot product = sum of d_k products
• Variances add: d_k
• Std = sqrt(d_k)

Now feed scores of eight, zero and minus eight into softmax. E to the minus eight is about zero point zero zero zero three three five, so the weights are about zero point nine nine nine seven, zero point zero zero zero three, and one in ten million. Softmax has collapsed to nearly one hot.

• Scores 8, 0, -8 give 0.9997, 0.0003, 0.0000001
• Nearly one-hot

A saturated softmax is a training disaster. The softmax derivative involves p times one minus p. With p at zero point nine nine nine seven, that is about three ten thousandths. The gradient almost vanishes, so the queries and keys barely learn. Dividing by root d k restores a standard deviation of one.

• Gradient of softmax: p(1 - p)
• p = 0.9997 gives about 0.0003

Let us verify it. Draw ten thousand random query and key pairs for each d k, and compute the standard deviation of the dot product. Without scaling it comes out near two, eight and twenty three for d k of four, sixty four and five hundred twelve. With the scaling, it stays near one every time. Scaling the scores by eight turns eight, zero, minus eight into one, zero, minus one, and softmax gives zero point six six five, zero point two four five, zero point zero nine zero.

• Simulate random q, k
• Unscaled std grows; scaled std stays near 1

Sometimes a token must not look at certain others. When a model predicts the next word, token two cannot see token three, or training would be cheating. And when sentences in a batch are padded to equal length, real tokens must ignore the padding. Both are handled by masks.

• No peeking at the future
• Ignore padding tokens

The trick is simple. Before softmax, replace blocked scores with minus infinity. E to the minus infinity is zero, so those positions get exactly zero weight, and the remaining weights renormalize. Token two with scores one point four one four, zero point seven zero seven, and minus infinity gives four point one one three, two point zero two eight, zero. Sum six point one four one.

• Set blocked scores to -inf
• e^(-inf) = 0

The weights for token two become zero point six seven zero and zero point three three zero, and zero for the future. A causal mask applies this across the whole matrix, keeping only the lower triangle. A padding mask blanks entire key columns for pad tokens. Masking after softmax would be wrong, because the rows would no longer sum to one.

• Causal: lower-triangular matrix
• Padding: mask whole key columns
• Weights: 0.670, 0.330, 0

A single attention head produces one set of weights per token. But a word can relate to others in several ways at once: grammar, nearby words, a distant reference. One softmax must compromise among them. Multi-head attention runs several smaller attentions in parallel, each free to specialize.

• One weighting pattern per token
• Language has many relations at once

With a model dimension of five hundred twelve and eight heads, each head works in a subspace of dimension sixty four. We project into queries, keys and values of width five twelve, then slice the last axis into eight chunks of sixty four. Each head runs scaled dot product attention on its own chunk, with its own scores and weights.

• d_model = 512, h = 8
• d_k = 512 / 8 = 64

After attention, each head outputs a vector of sixty four numbers per token. We concatenate the eight outputs back into width five twelve, then multiply by an output matrix W O of shape five twelve by five twelve. Without it, heads would never mix; W O lets the model combine what the heads found.

• Concat gives width 512
• W_O is 512 by 512
• Mixes information across heads

Shape tracking prevents most bugs. Input is batch B, time T, width five twelve. After projection it is the same. Reshape to B, T, 8, 64, then transpose to B, 8, T, 64. Scores are B, 8, T, T. Multiplying by V gives B, 8, T, 64. Transpose back, reshape to B, T, 512, and apply W O.

• (B,T,512) to (B,T,8,64)
• Transpose to (B,8,T,64)
• Scores: (B,8,T,T)

The four matrices W Q, W K, W V and W O each hold five twelve squared weights, about one point zero five million in total, ignoring biases. Notice that more heads do not add parameters, they only slice the width. Two classic bugs: reshaping straight to B, 8, T, 64 without transposing scrambles tokens, and applying softmax over the wrong axis. Softmax must run over the key axis.

• 4 × 512 × 512 ≈ 1.05M weights
• Mistake: reshape without transpose
• Mistake: wrong softmax axis

In NumPy, attention is five lines. Compute scores as Q matrix multiplied by K transpose, divided by the square root of d k. Add the mask, with minus infinity where blocked. Softmax each row, subtracting the row maximum first for numerical stability, as in lesson 15. Finally multiply the weights by V. Our hand example reproduces one point seven two, zero point seven two.

• scores = Q @ K.T / sqrt(d_k)
• Subtract the row max
• Apply mask, then softmax

PyTorch follows the same steps, batched. Use view and transpose to split heads, matmul for scores, masked fill with minus infinity for the causal mask, and softmax on the last dimension. Then transpose, reshape, and apply a linear layer for W O. Add shape assertions as you go. PyTorch also ships a fused function, scaled dot product attention, which you should use in production.

• Reshape, transpose, matmul
• masked_fill with -inf
• Check shapes with assert

Plot the weight matrix as a heatmap, with queries on the rows and keys on the columns. Each row sums to one. With a causal mask, the upper triangle is exactly zero. Plot all eight heads and you will often see different patterns: some focus on the previous token, some spread broadly. Heatmaps are useful for debugging, though they are not a full explanation of model behavior.

• Rows: queries; columns: keys
• Causal mask: lower triangle
• Different heads, different patterns

Now cost. Q times K transpose multiplies T by d with d by T, costing on the order of T squared times d. The score matrix itself has T squared entries per head. Double the sequence length and cost quadruples. At T equal to four thousand ninety six, one head holds almost seventeen million scores. That quadratic wall motivates the efficient attention methods we meet in lesson 58.

• T×T score matrix per head
• Double T, quadruple cost
• Memory grows with T²

Let us recap. Each token makes a query, key and value. Scores are dot products, softmax turns them into weights, and the output is a weighted sum of values. Dot product variance grows with d k, so we divide by its square root to avoid saturation. Masks use minus infinity before softmax. Multiple heads split the width, run in parallel, and recombine through W O. The cost is order T squared times d.

• QKV, scores, softmax, weighted sum
• Scale by 1/sqrt(d_k)
• Masks, heads, W_O

In the previous lesson we turned text into tokens and embeddings. Today we turned those embeddings into context. Next time, in lesson 56, we wrap attention inside the full transformer block, with residual connections, layer normalization and feed forward layers.

• Residuals, layer norm, MLP
• Subscribe for lesson 56