Multi-Head Attention - Transformer Architecture

Multi-Head Attention

What you'll build

The head-splitting core of multi-head attention: cut the model dimension into num_heads slices, run scaled dot-product attention on every slice in parallel, and concatenate the results.

Background

One attention head produces one weighted average per query, so it can only "look" one way at a time. Several smaller heads let the model attend to different things at once (for example, the previous token in one head and the subject of the sentence in another) for the same cost as one big head. In a full layer, Q, K, V come from learned projections X @ W_Q etc. and the concatenated heads go through an output projection W_O. Here the inputs are already projected and there is no output projection: the task is the reshape-attend-merge part, which is where the shape bugs live.

Formula

With d_k = d_model / num_heads, head h uses columns h*d_k : (h+1)*d_k:

head_h = softmax(Q_h @ K_h.T / sqrt(d_k)) @ V_h          # (seq_len, d_k)
output = concat(head_0, ..., head_{H-1}, axis=-1)        # (seq_len, d_model)

Vectorised: reshape (seq_len, d_model) -> (seq_len, H, d_k), transpose to (H, seq_len, d_k), do batched attention, then undo both steps.

Inputs

• Q, K, V: float arrays (seq_len, d_model)
• num_heads: int that divides d_model

Expected Return Format

output: array (seq_len, d_model), head h's result in columns h*d_k : (h+1)*d_k.

Edge cases and conventions

• Scale by sqrt(d_k) (the per-head width), not sqrt(d_model).
• num_heads = 1 is plain scaled dot-product attention.
• Stable softmax (subtract the row max).

Constraints

NumPy only. Use reshape/transpose and batched matmul; no loop over heads is needed.

Follow-up

• Add learned W_Q, W_K, W_V, W_O. How many parameters does the layer have? (Answer: 4 · d_model^2, independent of num_heads.)
• Why does reshape(seq_len, d_model) directly on the (H, seq_len, d_k) array give the wrong answer?

Input: Q = K = np.array([[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]), V = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]), num_heads = 2
Output: [[2.321, 3.321, 4.321, 5.321], [3.679, 4.679, 5.679, 6.679]]

Input: Q (4, 64), K (4, 64), V (4, 64), num_heads = 8
Output: output shape (4, 64)

import numpy as np

def multi_head_attention(Q, K, V, num_heads):
    """
    Multi-head attention on already-projected Q, K, V (no learned weights).

    Args:
        Q: Queries (seq_len, d_model)
        K: Keys (seq_len, d_model)
        V: Values (seq_len, d_model)
        num_heads: Number of attention heads (divides d_model)

    Returns:
        output: (seq_len, d_model), head h in columns h*d_k:(h+1)*d_k
    """
    seq_len, d_model = Q.shape
    assert d_model % num_heads == 0
    d_k = d_model // num_heads

    # Your code here
    pass