Multi-Head Attention
~18 mincode completion
Implement multi_head_attention(Q, K, V, num_heads).
Examples
Identity Q=K=V, one head
- Input
- multi_head_attention([[1, 0], [0, 1]], [[1, 0], [0, 1]], [[1, 0], [0, 1]], 1)
- Output
- [[0.66976, 0.33024], [0.33024, 0.66976]]
Hints
Hint 1
Subtract the row max before exponentiating to keep the result stable.
Hint 2
Watch for this: Forgetting sqrt(d k) scaling.
Requirements
Use stable softmax (subtract row max before exp).
Split by contiguous column slices.
Concatenate outputs along axis=1.
num_heads: Number of heads (assume d_model % num_heads == 0)Return Output of shape (seq_len, d_model)
Constraints
Allowed library: NumPy only
Time limit: 200 ms, Memory: 64 MB
Try similar problems(4)
Where this shows up
~18 min
••••••••••••••••
8 employers weight this skill
4 frontier labs, 3 AI product companies, 1 enterprise vendor. Top match scores 91.
Python
import numpy as np
def multi_head_attention(
Q: np.ndarray,
K: np.ndarray,
V: np.ndarray,
num_heads: int,
) -> np.ndarray:
"""
Compute multi-head scaled dot-product attention.
Args:
Q, K, V: Arrays of shape (seq_len, d_model)
num_heads: Number of heads (assume d_model % num_heads == 0)
Returns:
Output of shape (seq_len, d_model)
"""
# 1) Compute d_k = d_model // num_heads
# 2) For each head h:
# - Slice Q, K, V column blocks for this head
# - Compute scores = (Q_h @ K_h.T) / sqrt(d_k)
# - Apply stable row-wise softmax
# - Compute head output = weights @ V_h
# 3) Concatenate all head outputs along axis=1
# YOUR CODE HERE
pass