Multi-Head Attention

Assemble the full multi-head attention block on top of single head attention.

Given input X of shape (n, d_model) and projection matrices W_q, W_k, W_v each (d_model, d_model) plus an output projection W_o of the same shape:

1. Project: Q = X W_q, K = X W_k, V = X W_v.
2. Split each into num_heads heads along the feature axis. Head h takes columns h * d_head to (h + 1) * d_head, where d_head = d_model / num_heads.
3. Run scaled dot product attention inside each head, scaling by sqrt(d_head) and not sqrt(d_model).
4. Concatenate the head outputs back to (n, d_model) in head order.
5. Project once more: out = concat W_o.

Return (output, weights) where output is (n, d_model) and weights is (num_heads, n, n), holding each head's attention matrix.

The heads cost nothing extra. The projection is the same size either way, it is just reshaped so that each head attends in its own d_head-dimensional subspace. One head has to average all its relationships into a single distribution; eight heads can track syntax, position and coreference separately.

mask is None or a boolean array broadcastable to (n, n), applied identically in every head, where False blocks a position.

Function signature

def multi_head_attention(X, W_q, W_k, W_v, W_o, num_heads, mask=None):

Example

Input

X = np.array([[1.0, 0.0], [0.0, 1.0]])
I = np.eye(2)
multi_head_attention(X, I, I, I, I, num_heads=2)

Output

(array([[0.73105858, 0.5       ],
       [0.5       , 0.73105858]]), array([[[0.73105858, 0.26894142],
        [0.5       , 0.5       ]],

       [[0.5       , 0.5       ],
        [0.26894142, 0.73105858]]]))

Reasoning

With identity projections, head 0 sees only column 0 and head 1 sees only column 1. In head 0, token 0 scores [1, 0], so its weights are softmax([1, 0]) = [0.731, 0.269]. Token 1 scores [0, 0], so it attends equally. Head 1 is the same with the tokens swapped.

Constraints and notes

• Scale by sqrt(d_head), not sqrt(d_model).
• Head h takes columns h * d_head to (h + 1) * d_head.
• weights has shape (num_heads, n, n).