Many Heads, One Vector

A single attention pattern has to choose. Its weights sum to $1$, so "slept" attending to its subject means "slept" attending less to its adverb — even though both relationships are real. Multi-head attention buys its way out: several attention mechanisms run side by side, each with its own budget, and their outputs are joined back into one vector.

The projections are already done. queries, keys and values are each an $n \times d_{model}$ table, and the heads are carved out of the width:

• $d_{head} = d_{model} / h$, where $h$ is num_heads and divides $d_{model}$ exactly.
• Head $t$ owns the contiguous column slice $t \cdot d_{head}$ through $(t+1) \cdot d_{head} - 1$ of all three tables, and never looks outside it.
• Within its slice, head $t$ runs ordinary self-attention over the whole sequence: score each query against every key, divide by $\sqrt{d_{head}}$, softmax across the row, then blend that head's values by those weights.
• Every head produces an $n \times d_{head}$ output. Concatenate them back in head order into an $n \times d_{model}$ table, then multiply that table on the right by the output matrix $W_O$, of shape $d_{model} \times d_{model}$: entry $[i][c]$ of the result is row $i$ of the joined table dotted with column $c$ of $W_O$.

Task: write multi_head_attention(queries, keys, values, num_heads, w_out) returning the final $n \times d_{model}$ table, every entry rounded to 4 decimal places. Round once, at the very end.

• Row i of the result is the output for token i.
• $h$ may be $1$ (plain single-head attention) and may be as large as $d_{model}$ (every head one column wide, so $\sqrt{d_{head}} = 1$).

The heads are narrow, which is why eight of them cost roughly what one wide one costs. It is also why the scale is $\sqrt{d_{head}}$: a dot product inside a head sums $d_{head}$ terms, and that count is what decides how large the raw scores grow.