Multi-Head Attention

Your program will run several smaller attentions side by side and recombine them. The single-head attention from the last lesson is given to you; the new work is cutting the vectors into heads, running them, and stitching the answers back into one matrix.

Why one head is not enough

One softmax produces one weighting per token. But a token often needs several things at once: the noun it refers to, the verb it belongs to, the previous word for rhythm. A single average over all of that blurs it. Multi-head attention gives each of H heads its own, smaller view, so each can learn a different matching rule.

The shapes

The model width is D and there are H heads, so each head works in d_k = D / H dimensions (GPT-2 small: 768 / 12 = 64). The projected matrices Q, K, V are T x D. Head h owns columns h*d_k up to h*d_k + d_k - 1 of each of them.

        columns:  0 .. dk-1 | dk .. 2dk-1 | ...   | (H-1)dk .. D-1
Q, K, V   (T x D)   head 0  |   head 1    | ...   |   head H-1

per head:   out_h = attention(Q_h, K_h, V_h)           T x d_k
join:       concat = [ out_0 | out_1 | ... | out_(H-1) ]   T x D
project:    result = concat @ W_O                        T x D

The heads do not see each other's columns. The cost matches one full-width head, because H attentions of width D/H have the same total number of multiplications as one of width D.

Worked trace

T = 2, D = 4, H = 2, so d_k = 2. Take Q = K = [[1,0,0,2],[0,1,2,0]] and V = [[1,2,3,4],[5,6,7,8]].

head 0 (columns 0-1):  Q_0 = K_0 = [[1,0],[0,1]]   V_0 = [[1,2],[5,6]]
  scores / sqrt(2) = [[0.7071, 0], [0, 0.7071]]  ->  weights [[0.6698, 0.3302], [0.3302, 0.6698]]
  out_0 = [[2.3210, 3.3210], [3.6790, 4.6790]]

head 1 (columns 2-3):  Q_1 = K_1 = [[0,2],[2,0]]   V_1 = [[3,4],[7,8]]
  scores / sqrt(2) = [[2.8284, 0], [0, 2.8284]]  ->  weights [[0.9442, 0.0558], [0.0558, 0.9442]]
  out_1 = [[3.2232, 4.2232], [6.7768, 7.7768]]

concat row 0 = [2.3210, 3.3210, 3.2232, 4.2232]

The two heads weight the same two tokens very differently (67% against 94% on the token itself), which is the point: each head's scores come only from its own two columns. Attention over all four columns at once would give one set of weights, 92% and 8%.

The output projection

After the join, every head's answer still sits in its own columns. W_O is a learned D x D matrix that mixes them: result[t][c] = sum over j of concat[t][j] * W_O[j][c]. It also guarantees the output has the same shape as the input, which the residual connection in lesson 9 requires.

Parameters

Four D x D matrices: W_Q, W_K, W_V and W_O, so 4 D^2 numbers per layer whatever H is. Modern models trim the K and V side: multi-query and grouped-query attention let several query heads share one K/V head, which shrinks the cache that generation has to keep.

Common mistakes

• Scaling by sqrt(D) instead of sqrt(d_k). Each head's dot products have d_k terms, so that is the right size to divide by.
• Splitting by stride (head h takes columns h, h+H, h+2H, ...) instead of contiguous blocks. The shapes still fit; the numbers are wrong.
• Multiplying by W_O with rows and columns swapped. A non-symmetric W_O shows it.
• Letting heads share one softmax, or averaging the head outputs instead of concatenating them.

Your exercise

Write multi_head: slice, attend, join, project. The starter parses the input and includes the single-head attention from the previous lesson.