Multi-Head Attention 🎯 (100/100!)

v1 SDPA 는 하나의 attention head. Transformer 의 실제 레이어는 Multi-Head:

$\text{MHA}(x) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W_o$

$\text{head}_h = \text{SDPA}(Q_h, K_h, V_h)$

where $Q_h, K_h, V_h$ 는 $x$ 를 각각 $W_q, W_k, W_v$ 로 projection 후 $H$ 개 head 로 분할한 slice.

왜 여러 head 가 필요한가

• 하나의 attention 은 한 가지 관계 에만 집중 (예: 문법적 연결).
• 여러 head = 여러 부분 공간 에서 서로 다른 관계 학습 (구문·의미·위치 등).
• 같은 $d_{\text{model}}$ 을 $H$ 등분하므로 파라미터 수는 동일 (projection 크기만 분할).

구현 순서

1. $Q = xW_q$, $K = xW_k$, $V = xW_v$ — 모두 (L, d_model).
2. 각각 (L, H, d_head) 로 reshape → (H, L, d_head) 로 transpose.
3. 각 head 별로: $\text{score} = Q_h K_h^T / \sqrt{d_{\text{head}}}$.
4. mask 적용 ((L, L) bool 을 head 축으로 broadcast), softmax, × $V_h$.
5. head 결과 concat (L, d_model) → $W_o$.

d_model % num_heads == 0 가정.

과제

함수 multi_head_attention(x, W_q, W_k, W_v, W_o, num_heads, mask=None) 를 완성하세요.

• x (L, d_model), W_q/W_k/W_v/W_o (d_model, d_model).
• num_heads 로 분할.
• 반환 (L, d_model).

테스트 케이스