Problems
Loading...

05/08Multi-Head Attention

Hard

Implement projected multi-head attention.

headi=softmax⁡ ⁣(QWiQ(KWiK)Tdk)VWiVhead_i = \operatorname{softmax}\!\left(\frac{QW_i^Q(KW_i^K)^{\mathsf T}}{\sqrt{d_k}}\right)VW_i^V MHA⁡(Q,K,V)=Concat⁡(head1,…,headh)WO\operatorname{MHA}(Q,K,V) = \operatorname{Concat}(head_1,\ldots,head_h)W^O

Here, h is num_heads and (d_k=d_{model}/h). The supplied projection matrices combine every head in one model-width matrix. Return a NumPy float64 array with shape (batch, query_length, d_model).

Loading visualization...

Examples

Input: Q, K, V: shape (1,3,4), W_q, W_k, W_v, W_o: shape (4,4), num_heads = 2

Output: [[[-0.038, -0.0259, -0.0355, -0.1], [-0.0368, -0.0275, -0.0348, -0.1016], [-0.039, -0.0249, -0.0349, -0.0975]]]

Explanation: Q, K, V and weight matrices are concrete arrays. Each of the 2 heads operates on d_k = 4/2 = 2 dimensions. After projecting, splitting into heads, applying scaled dot-product attention, and concatenating, the output projection produces the final result.

Input: Q, K, V: shape (1,2,4), W_q, W_k, W_v, W_o: shape (4,4), num_heads = 1

Output: [[[-0.011, -0.1359, -0.1294, 0.0896], [-0.0125, -0.1353, -0.1289, 0.0885]]]

Input: Q, K, V: shape (2,2,4), W_q, W_k, W_v, W_o: shape (4,4), num_heads = 2

Output: [[[-0.0815, -0.0442, 0.0259, -0.1741], [-0.0721, -0.0332, 0.0363, -0.1683]], [[0.0471, 0.0347, -0.0182, 0.0653], [0.0538, 0.0437, -0.0111, 0.0725]]]

Hint 1

Reshape to batch, sequence, heads, head_width, then transpose the head axis forward.

Hint 2

Transpose only the final two dimensions of projected K.

Hint 3

Transpose heads back before reshaping and applying W_o.

Requirements

  • Project Q, K, and V before splitting heads.
  • Split the final width evenly across num_heads.
  • Scale each head's scores by the square root of its width.
  • Apply softmax over key positions.
  • Concatenate heads and apply W_o.

Constraints

  • Inputs and weights are compatible NumPy float64 arrays.
  • d_model is divisible by num_heads.
  • K and V have the same sequence length.
  • Valid inputs are guaranteed.
Prev5 / 8Next