multi-head attention — split, attend in parallel, concatenate
X (N, d_model)
@ W_q
@ W_k
@ W_v
split into 4 heads
Q (N, d_model) → reshape (4, N, d_head)
Q_1
Q_2
Q_3
Q_4
K_1
K_2
K_3
K_4
V_1
V_2
V_3
V_4
softmax(Q_h K_h^T / √d_head) V_h — batched matmul, one op
head 1
subject-verb
head 2
positional
head 3
copy / induction
head 4
named entities
concat → (N, d_model)
@ W_o (d_model, d_model)
output (N, d_model)
W_o is where the heads mix. until W_o, every head lives in its own subspace.