Multi-Head Attention

Task

Implement multi_head_attention(query: list[list[float]], key: list[list[float]], value: list[list[float]], w_query: list[list[float]], w_key: list[list[float]], w_value: list[list[float]], w_output: list[list[float]], num_heads: int, mask: list[list[bool]]) -> list[list[float]]. Apply learned query, key, and value projections, run scaled dot-product attention independently in multiple heads, concatenate the head outputs, and apply the output projection.

Requirements

• query, key, and value have model width d_model; every projection matrix has shape d_model by d_model.
• d_model must divide evenly by num_heads; each head has width d_head = d_model / num_heads.
• Project query, key, and value before splitting their last dimension into contiguous heads.
• Within each head, divide query-key dot products by sqrt(d_head).
• A True mask entry is valid and a False entry is blocked; the same query-key mask applies to every head.
• Use row-wise stable softmax over valid keys and return zeros for every fully masked query head.
• Concatenate heads in their original order and multiply by w_output.
• Preserve an empty query matrix by returning an empty list.

Example

query = [[1.0, 0.0, 0.0, 1.0]]
key = [[1.0, 0.0, 1.0, 0.0], [0.0, 1.0, 0.0, 1.0]]
value = [[10.0, 0.0, 1.0, 2.0], [0.0, 20.0, 3.0, 4.0]]
identity = [[1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 1.0]]

multi_head_attention(
    query, key, value,
    identity, identity, identity, identity,
    2, [[True, True]],
)
# [[6.69761549, 6.60476901, 2.33952310, 3.33952310]]