3.2 Self-Attention (Q, K, V)

Scaled dots, softmax over keys, mix values. Multi-head, then a causal mask.

These notes match the lecture slides. Use (slides) on the course hub for the deck.

Self-attention is a weighted sum of values, with weights from a comparison of queries to keys. Every token plays all three roles. That is the calculation, and Lab 3. By the end you should be able to write the scaled-dot formula, fill a $2 \times 2$ softmax table, and mask the future with $- \infty$ , not zero.

1. What self-attention is

Self-attention is the calculation inside a transformer layer. A weighted sum of values, with weights from queries compared to keys. Every token plays all three roles.

For a sequence matrix $X \in ℝ^{T \times d}$ ,

Q = X W_{Q}, K = X W_{K}, V = X W_{V} .

Query: what this token is looking for. Key: what another token looks like to be found. Value: what you actually add into the representation.

Row $t$ of $Q K^{⊤}$ is “how much token $t$ wants each position.” Softmax makes that a distribution. Then you mix values, not keys. Mixing keys is a common lab bug.

Query, key, value, then a mix of values

Self versus cross. Self-attention: $Q, K, V$ all from the same sequence. Cross-attention (encoder–decoder, later BLIP): queries from one sequence, keys and values from another. Same formula.

2. Why we use it

RNNs mix left-to-right through $h_{t}$ . Self-attention mixes any pair in one layer, in parallel. Note 3.1 argued for path length 1; this note is the arithmetic that implements it.

Why $\sqrt{d_{k}}$ . If the coordinates of $q$ and $k$ are roughly independent with variance 1, the dot product has variance $d_{k}$ . Large $d_{k}$ pushes softmax toward a one-hot. Dividing by $\sqrt{d_{k}}$ puts the logits back in a range where softmax still mixes. Skipping the scale, or dividing by $d_{k}$ instead, is the other common exam slip.

Several heads in parallel let different slices of the channels track syntax, a name, or a comma. One head is enough to compute by hand in Lab 3; multi-head is why the layer is not a single average.

3. Architecture

Three matrices map $X$ to $Q, K, V$ . Scores are $Q K^{⊤} / \sqrt{d_{k}}$ , softmax over keys, then mix $V$ :

A = s o f t m a x (\frac{Q K^{⊤}}{\sqrt{d_{k}}}), A t t e n t i o n (Q, K, V) = A V .

Multi-head: several $(W_{Q}, W_{K}, W_{V})$ in parallel, each on a slice of the channels (head dimension $d_{k} = d / h$ ). Concatenate the head outputs, then a linear $W_{O}$ . Concat is not a residual, and $W_{O}$ is part of attention’s parameter count in note 3.3.

Two heads, concatenate, project

Lab 3 computes one head by hand. You do not need eight heads on the board.

For next-token prediction you mask future keys (set scores to $- \infty$ before softmax). Otherwise the model cheats. BERT does not mask the future; it masks random tokens and reconstructs them (note 3.4). Those are two different uses of the word mask.

Zero is the wrong fill-in: $s o f t m a x ([1, 0])$ still puts mass on the “zero.” $- \infty$ becomes probability $0$ .

Causal mask zeros out the future

4. How it works, step by step

Form $Q, K, V$ from $X$ (or, in a tiny drill, start from scores $S = Q K^{⊤}$ already).
Scale: divide by $\sqrt{d_{k}}$ .
If this is a language model, set illegal future scores to $- \infty$ .
Softmax over keys (last dim / each row), not over columns.
Multiply the resulting distribution by $V$ . Output row $t$ is a mix of values.

Row $t$ of $A$ is a distribution over the sequence. Large mass on a name is “this pronoun just looked at that name.”

Attention weights for one head

If two keys are equal, the softmax splits mass equally between them (given the same query). That is not a tie-break bug; it is the definition.

Lab 3 is this procedure at $T = 4$ : one head by hand, then a causal mask. Do not start that notebook until the $2 \times 2$ arithmetic below has been done once on paper.

5. Mathematical formulas

The layer:

Attention (Q, K, V) = softmax (Q K^{⊤} / \sqrt{d_{k}}) V,

with $Q = X W_{Q}$ , $K = X W_{K}$ , $V = X W_{V}$ .

Causal mask: before the softmax, illegal (future) entries of $Q K^{⊤} / \sqrt{d_{k}}$ are set to $- \infty$ , so those probabilities become 0.

Multi-head: $h$ heads, each of width $d_{k} = d / h$ ; concatenate, then $W_{O}$ . If two heads each emit a 2-D vector and the model width is $d = 4$ , $W_{O}$ maps the concatenated 4-D vector back to width 4.

6. Positive points and negative points

Positive.

Parallel over tokens, path length 1: any pair can interact in one layer.
Multi-head specialization: different heads can track syntax, a name, or a comma.
The same formula covers self-attention and cross-attention; only which sequence supplies $Q$ versus $K, V$ changes.
A heatmap of $A$ is a debug tool (Lab 3), not just a picture.

Negative.

$O (T^{2})$ scores per head per layer, as note 3.1 already billed.
Mixing keys instead of values is a common bug; the softmax weights multiply $V$ .
No positions until note 3.3: attention is a set operation unless you add order.
Dividing by $d_{k}$ instead of $\sqrt{d_{k}}$ , or skipping the scale, peaks or flattens softmax for the wrong reason.
Masking with $0$ instead of $- \infty$ still puts mass on the illegal key.
Softmax over the wrong axis (columns instead of keys) silently permutes the mix.

When not to. If the task is next-token prediction, bidirectional self-attention is cheating. Use a causal mask. BERT’s token mask is a different objective (note 3.4), not a substitute for that triangular mask.

7. Teaching this note

~18 minutes. Define Q, K, V in one sentence each, write the scaled-dot formula, grind the $2 \times 2$ table (including $\sqrt{d_{k}}$ ), then the causal mask. Play the QKV and softmax stretch of 3Blue1Brown attention (about 4:00–16:00). Lab 3 is $T = 4$ ; do not start that notebook until this arithmetic is done once by hand.

8. Worked example

$T = 2$ , $d_{k} = 2$ , skip $W_{Q}$ and take scores $S = Q K^{⊤}$ already:

S = [\begin{matrix} 2 & 0 \\ 0 & 2 \end{matrix}], \frac{S}{\sqrt{2}} \approx [\begin{matrix} 1.41 & 0 \\ 0 & 1.41 \end{matrix}] .

Row 1 softmax: $e^{1.41} \approx 4.10$ , $e^{0} = 1$ , so $[4.10, 1] / 5.10 \approx [0.80, 0.20]$ .

Row 2 is the same on the diagonal: $A \approx [\begin{matrix} 0.80 & 0.20 \\ 0.20 & 0.80 \end{matrix}]$ .

Scaled scores to a mix of values

If $V = I$ , then $A V = A$ . Each output is 80% “self” and 20% “the other token.”

Causal mask on row 1: set the future score to $- \infty$ . Then row 1 softmax is $[1, 0]$ : token 1 cannot look at token 2.

Without the $\sqrt{2}$ , row 1 would be $s o f t m a x ([2, 0]) \approx [0.88, 0.12]$ : more peaked. That is the scale’s job.

9. Where students get stuck

Softmax over the wrong axis (columns instead of keys / last dim).
Dividing by $d_{k}$ instead of $\sqrt{d_{k}}$ , or skipping the scale.
Masking with $0$ instead of $- \infty$ .
Mixing keys instead of values.

10. Video

Watch 3Blue1Brown: Attention in transformers, visually explained.

Pause when queries and keys meet as a grid of dots, when softmax turns a row into a distribution, and when values are mixed. That grid is Lab 3’s heatmap.

11. Practice

If two tokens have identical keys, what does the softmax do to their weights?
Why divide by $\sqrt{d_{k}}$ ?
Write $A$ for $T = 2$ with $Q = K = V = I$ (identity), no mask. What is $A V$ ?
Softmax the row $[0, 2]$ . Then softmax $[0, 2] / \sqrt{2}$ . Did the larger entry gain or lose probability after scaling?
Scores $[\begin{matrix} 1 & 3 \\ 0 & 1 \end{matrix}]$ , causal mask (upper triangle $- \infty$ ). Write the masked score matrix, then $A$ (row-wise softmax).
Self-attention versus cross-attention: which matrices come from which sequence?
Two heads each emit a 2-D vector. After concat, what width does $W_{O}$ map back to if the model width is $d = 4$ ?

Ahmad Mousavi