Lesson 80: Scaled Dot-Product Attention from Scratch

Discussion prompt

Before we open Lesson 80: Scaled Dot-Product Attention from Scratch: without looking back, what was the main idea of Scaled Dot-Product Attention, and what could you do by the end of it that you could not do before?

Hint: One sentence for the idea, one for the skill. If the second one is blank, that is the part to revisit.

Discussion prompt

Attention answers: for each query token, how much should I attend to each key token? The answer is a weighted sum over the values.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Sort into buckets

Sort these by shape, from memory, without looking back at QKV shape contract. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

(batch, seq, d_k)

Q (query); K (key)

(batch, seq, d_v)

V (value); output

g1

shape is "(batch, seq, d_k)" for Q (query), K (key) — that is what the table on "QKV shape contract" records, and it is the single property separating this group from the rest.

g2

shape is "(batch, seq, d_v)" for V (value), output — that is what the table on "QKV shape contract" records, and it is the single property separating this group from the rest.

Discussion prompt

Explain Why divide by sqrt(d_k)? by analogy to something with no Machine Learning in it at all — a queue, a recipe, a map, a bank balance, whatever fits. Then say where your analogy breaks.

Hint: An analogy that never breaks is not an analogy, it is the same idea wearing a hat. Find the seam — that is the part that is actually new.

Predict first

Build scaled_dot_product_attention step by step. Start: Q @ K^T then scale.

Commit before you compute: what does Implement attention: shapes and scores come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: raw variance = 6.5991 → scaled variance = 0.8249 (near 1)

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Dividing by sqrt(8)=2.828 normalises the dot-product distribution.

Sort into buckets

Sort these by shape, from memory, without looking back at Implement attention: shapes and scores. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

(2, 5, 5)

Q @ K^T (raw scores); / sqrt(d_k=8); softmax(scores, dim=-1)

(2, 5, 8)

weights @ V

g1

shape is "(2, 5, 5)" for Q @ K^T (raw scores), / sqrt(d_k=8), softmax(scores, dim=-1) — that is what the table on "Implement attention: shapes and scores" records, and it is the single property separating this group from the rest.

g2

shape is "(2, 5, 8)" for weights @ V — that is what the table on "Implement attention: shapes and scores" records, and it is the single property separating this group from the rest.

Discussion prompt

Add softmax over the key dimension (dim=-1) and multiply by V. Verify the output shape.

What do you need to know — or decide — before the first line can be written? List everything the problem has to hand you.

Hint: Anything you would have to invent to get started is a thing the problem must supply.

Comparison matrix

From Softmax → context output: refill the attention weight (query 0) column from what you know. The rest of the table is as it appeared.

Comparison matrix

From Why masking is necessary: every row here is a choice with a cost. Fill the what it blocks column, then say which row you would actually pick and what you give up for it.

Step through it

Step through Causal mask: lower-triangular boolean one row at a time. What is driving the change, and what would the row after the last one be?

1. Step 1: row (query) is 0
2. Step 2: row (query) is 1
3. Step 3: row (query) is 2
4. Step 4: row (query) is 3
5. Step 5: row (query) is 4

Predict first

The table runs: 0 (causal) | 1.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 · 2 (causal) | 0.2639 | 0.1116 | 0.6245 | 0.0000 | 0.0000

In Apply causal mask with masked_fill, given the rows so far: what is the next one — the row where query pos is 2 (unmasked)?

Correct: 2 (unmasked) | 0.1005 | 0.0425 | 0.2378 | 0.0937 | 0.5256

Why: The relationship between the columns, not the individual numbers, is what generates the next row. After -inf fill, exp(-inf)=0 exactly.

Step through it

Step through Apply causal mask with masked_fill one row at a time. What is driving the change, and what would the row after the last one be?

1. Step 1: query pos is 0 (causal)
2. Step 2: query pos is 2 (causal)
3. Step 3: query pos is 2 (unmasked)

Predict first

A student writes this, and it looks reasonable:

Compute weights = F.softmax(scores, dim=-1), then zero out the forbidden positions: weights[~mask] = 0.0.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: Zeroing AFTER softmax leaves weights [0.091, 0, 0, 0, 0] — they no longer sum to 1.

Set blocked scores to -inf BEFORE softmax: scores.masked_fill(~mask, float('-inf')).

Why: Zeroing AFTER softmax leaves weights [0.091, 0, 0, 0, 0] — they no longer sum to 1. The context vector is scaled by 0.091, not 1. Downstream layers receive wrong magnitudes, gradients are corrupted.

Discussion prompt

The rule this trap just fixed:

Set blocked scores to -inf BEFORE softmax: scores.masked_fill(~mask, float('-inf')).

Now break it on purpose. Build a case that violates it and follow the consequences until something visibly fails. Where does the failure first show up — and would you have noticed it if you had not been looking?

Hint: The dangerous rules are the ones whose violation still produces an answer. If yours fails loudly, try to find one that fails quietly.

Discussion prompt

Explain FP16 overflow in naive softmax to a student a year behind you. No notation, no jargon they have not met — and it still has to be true.

Hint: If your explanation needs a symbol they have never seen, you are describing the notation rather than the idea.

Discussion prompt

Explain The max-subtraction fix by analogy to something with no Machine Learning in it at all — a queue, a recipe, a map, a bank balance, whatever fits. Then say where your analogy breaks.

Hint: An analogy that never breaks is not an analogy, it is the same idea wearing a hat. Find the seam — that is the part that is actually new.

Predict first

Assemble the complete function with optional mask and the safe softmax path.

Commit before you compute: what does Full scratch implementation + stability come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: match = True, max abs diff = 2.38e-07

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Our implementation is numerically identical to PyTorch's fused kernel (within float32 rounding).

Discussion prompt

Work backwards. The example finished here:

match = True, max abs diff = 2.38e-07

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Comparison matrix

From Computational cost: O(T² · d_k): refill the attn matrix entries column from what you know. The rest of the table is as it appeared.

Discussion prompt

Walk every number by hand with T=3, d_k=4, d_v=2. Fixed Q3, K3, V3 (no randomness — exact reproducible values).

What do you need to know — or decide — before the first line can be written? List everything the problem has to hand you.

Hint: Anything you would have to invent to get started is a thing the problem must supply.

Predict first

Apply the causal mask to the same Q3/K3/V3 and trace how query 1's weights redistribute.

Commit before you compute: what does Causal mask on the 3-token example come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: q0 attends only to k0 (weight=1.0); q1 attends to k0 (0.5622) and k1 (0.4378); q2 unaffected (same as unmasked)

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. q2 is the last token — it can see all 3 keys, so masking changes nothing for it.

Discussion prompt

Work backwards. The example finished here:

q0 attends only to k0 (weight=1.0); q1 attends to k0 (0.5622) and k1 (0.4378); q2 unaffected (same as unmasked)

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Put in order

These are the steps of The scaled dot-product attention recipe, scrambled. Put them back in order before the next slide shows you.

1. Shapes: Q/K (batch, seq, d_k), V (batch, seq, d_v) → output (batch, seq, d_v)
2. Scores: Q @ K.transpose(-2,-1) → (batch, seq, seq); divide by sqrt(d_k)
3. Mask (if any): scores.masked_fill(~mask, float('-inf')) — always BEFORE softmax
4. Weights: F.softmax(scores, dim=-1) — safe against overflow, rows sum to 1
5. Output: weights @ V — weighted sum of values; verify against F.scaled_dot_product_attention

Why: This is the order the recipe itself gives. Recalling the sequence without the slide in front of you is the difference between recognising the method and being able to run it — most of what goes wrong in practice is a step done out of turn.

Discussion prompt

The scaled dot-product attention recipe works on the cases you have just seen. Push it to the edge: what is the most degenerate input it still handles — empty, zero, one item, everything equal — and what is the first case where it stops being true? Name the case, not just "it breaks".

Hint: Try the smallest legal input, then the largest, then the one where two things collide. Methods are specified at their edges; the middle takes care of itself.

Eliminate the wrong options

Q has shape (4, 10, 64), K has shape (4, 10, 64), V has shape (4, 10, 128). What is the shape of the attention output?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. (4, 10, 128)
• B. (4, 10, 64)
• C. (4, 64, 128)
• D. (4, 10, 10)

Survives elimination: A

Why: Output shape = (batch, seq, d_v). The attention-weight matrix is (4,10,10); multiplying by V of shape (4,10,128) gives (4,10,128). d_v=128 sets the output width, not d_k=64.

Check your understanding

Q has shape (4, 10, 64), K has shape (4, 10, 64), V has shape (4, 10, 128). What is the shape of the attention output?

• A. (4, 10, 128) (correct)
• B. (4, 10, 64)
• C. (4, 64, 128)
• D. (4, 10, 10)

Answer: A

Why: Output shape = (batch, seq, d_v). The attention-weight matrix is (4,10,10); multiplying by V of shape (4,10,128) gives (4,10,128). d_v=128 sets the output width, not d_k=64.

Why B tempts people

(4,10,64) confuses d_k (the key/query dimension used for dot products) with d_v (the value dimension that determines output width). They are equal in single-head standard attention but differ in many architectures.

Why C tempts people

(4,64,128) transposes the seq and d_k axes — there is no operation in SDPA that produces a (d_k, d_v) output. The seq dimension is always preserved.

Why D tempts people

(4,10,10) is the shape of the intermediate attention-weight matrix (before multiplying by V) — it is not the final output.

Predict first

For d_k=64 (a typical transformer head size), what is the correct scale divisor, and why 1/sqrt(d_k) specifically?

Answer it in your own words, now, with nothing to choose from. The options are on the next slide — and picking the right one off a list is an easier skill than producing it.

Correct: 8.0 — normalises dot-product variance to ~1 for unit-normal Q, K

Why: For q, k ~ N(0,1), Var(q·k) = d_k, so the standard deviation is sqrt(d_k) = sqrt(64) = 8. Dividing by sqrt(d_k) gives Var = 1. Verified: raw variance 6.60 → scaled variance 0.82 at d_k=8.

Check your understanding

For d_k=64 (a typical transformer head size), what is the correct scale divisor, and why 1/sqrt(d_k) specifically?

• A. 8.0 — normalises dot-product variance to ~1 for unit-normal Q, K (correct)
• B. 64.0 — divides by d_k so each score is a per-element average
• C. 0.5 — a fixed constant independent of d_k
• D. 1.0 — no scaling is needed; softmax is scale-invariant

Answer: A

Why: For q, k ~ N(0,1), Var(q·k) = d_k, so the standard deviation is sqrt(d_k) = sqrt(64) = 8. Dividing by sqrt(d_k) gives Var = 1. Verified: raw variance 6.60 → scaled variance 0.82 at d_k=8.

Why B tempts people

Dividing by d_k (not sqrt(d_k)) over-shrinks the scores: variance becomes 1/d_k instead of 1, pushing softmax toward uniform distribution and reducing expressiveness.

Why C tempts people

A fixed constant like 0.5 does not adapt to d_k. As d_k grows the scores still saturate softmax — the scaling must be 1/sqrt(d_k) to maintain unit variance regardless of head size.

Why D tempts people

Softmax is NOT scale-invariant in practice: large inputs saturate it (gradients near zero). The variance of dot products grows with d_k, so large d_k without scaling collapses attention weights to near-one-hot spikes.

Eliminate the wrong options

A student writes: weights = F.softmax(scores, dim=-1); weights[~causal_mask] = 0.0. What is wrong?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. The weights no longer sum to 1 — the context vector has incorrect magnitude
• B. Setting to 0.0 is identical to setting the score to -inf before softmax
• C. causal_mask should be the upper triangle, not the lower triangle
• D. Nothing — zeroing after softmax and -inf before softmax are mathematically equivalent

Survives elimination: A

Why: After softmax, a row like [0.09, 0.02, 0.03, 0.18, 0.68] sums to 1.0. Zeroing the last 4 leaves [0.09, 0, 0, 0, 0] — sum = 0.09, not 1. The context vector is scaled by 0.09 instead of 1, corrupting downstream projections and gradients. The -inf-before-softmax trick avoids this because exp(-inf)=0 before normalisation.

Check your understanding

A student writes: weights = F.softmax(scores, dim=-1); weights[~causal_mask] = 0.0. What is wrong?

• A. The weights no longer sum to 1 — the context vector has incorrect magnitude (correct)
• B. Setting to 0.0 is identical to setting the score to -inf before softmax
• C. causal_mask should be the upper triangle, not the lower triangle
• D. Nothing — zeroing after softmax and -inf before softmax are mathematically equivalent

Answer: A

Why: After softmax, a row like [0.09, 0.02, 0.03, 0.18, 0.68] sums to 1.0. Zeroing the last 4 leaves [0.09, 0, 0, 0, 0] — sum = 0.09, not 1. The context vector is scaled by 0.09 instead of 1, corrupting downstream projections and gradients. The -inf-before-softmax trick avoids this because exp(-inf)=0 before normalisation.

Why B tempts people

They are NOT equivalent: -inf before softmax sets exp to 0 before the denominator is computed, so normalisation corrects for it. Zeroing after softmax skips the renormalisation step entirely.

Why C tempts people

The lower-triangular mask is correct for causal attention — it allows each position to attend to itself and past positions (i.e., j ≤ i). The upper triangle is what gets blocked.

Why D tempts people

They are not equivalent — see answer A. This is one of the most common attention implementation bugs.

Discussion prompt

Implement scaled_dot_product_attention(Q, K, V, mask=None), test on the 3-token example, add a causal mask, and verify against F.scaled_dot_product_attention.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Step through it

Step through Milestone 1 & 2 — scores and weights one row at a time. What is driving the change, and what would the row after the last one be?

1. Step 1: query is q0
2. Step 2: query is q1
3. Step 3: query is q2

Comparison matrix

From Milestone 3 & 4 — causal mask + verification: every row here is a choice with a cost. Fill the result column, then say which row you would actually pick and what you give up for it.

Comparison matrix

From The full program: refill the match vs F.scaled_dot_product_attention column from what you know. The rest of the table is as it appeared.

Draw it

One page, no notation unless you need it: draw how these connect — QKV shapes and the attention formula · Masking: causal and padding · Numerical stability in FP16 · Complexity and the 3-token trace · Your turn: implement it. Put an arrow wherever one of them is what makes another possible, and label the arrow with why.