Neural Mastery
← Practice
Multi-Head Self-Attention From Scratchhard

Multi-Head Self-Attention From Scratch

hardDeep Learning⏱ 20–25 min
⚡ Libraries allowed · Pure Python earns +10 bonus XP
🎯 Mission
Implement multi-head self-attention projections, scaled dot-product attention, and head concatenation from scratch in pure NumPy.

Task

Implement `multi_head_attention(Q, K, V, d_model, num_heads)`. Libraries (NumPy) are allowed!

Function Signature

multi_head_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray, d_model: int, num_heads: int) -> np.ndarray

Examples

Example 1: Output Shape Test
Input: {"Q":[[1,0],[0,1]],"K":[[1,0],[0,1]],"V":[[1,2],[3,4]],"d_model":2,"num_heads":1}
Output: [[2,3],[2,3]]

Constraints

  • d_model must be evenly divisible by num_heads.
  • Must compute scaled dot-product attention softmax(Q K^T / sqrt(d_k)) V per head.
  • Output shape must strictly match (seq_len, d_model).
Python3Saved ✓
import numpy as np

def multi_head_attention(Q, K, V, d_model, num_heads):
"""
Q, K, V: 2D arrays of shape (seq_len, d_model)
d_model: total embedding dimension
num_heads: number of attention heads
Returns: output array of shape (seq_len, d_model)
"""
# Your implementation here
pass

Case 1: Output Shape Test
Input: {"Q":[[1,0],[0,1]],"K":[[1,0],[0,1]],"V":[[1,2],[3,4]],"d_model":2,"num_heads":1}
Expected: [[2,3],[2,3]]