Multi-head attention

Implementasikan causal multi-head self-attention:

$Q = XW_q,\ K = XW_k,\ V = XW_v$ , semuanya (T, D).
Pecah dimensi D menjadi n_heads head berukuran hd = D / n_heads.
Untuk setiap head, hitung causal attention dengan skala $1/\sqrt{hd}$ .
Gabungkan kembali semua head menjadi (T, D), lalu kalikan dengan $W_o$ .

Head ke-h memakai kolom h*hd : (h+1)*hd dari Q, K, dan V.

Kenapa banyak head

Satu head attention hanya menghasilkan satu pola bobot per token, padahal satu token sering butuh beberapa jenis informasi sekaligus: subjek kalimatnya, kata sebelumnya, tanda kurung yang belum ditutup. Dengan $H$ head, model bisa memperhatikan $H$ hal berbeda secara paralel. Biayanya hampir sama dengan satu head besar, karena setiap head hanya berdimensi $D/H$ .

Trik reshape

Tidak perlu loop per head. Cukup:

(T, D) --reshape--> (T, H, hd) --transpose--> (H, T, hd)

Sekarang H berperan sebagai dimensi batch, dan matmul berbatch dari soal P0 mengerjakan semua head sekaligus. Kebalikannya: transpose(1, 0, 2).reshape(T, D).

Kalau ada dimensi batch B juga, bentuknya menjadi (B, H, T, hd). Itulah bentuk yang akan kamu lihat di kode GPT mana pun.

W_o

Setelah digabung, setiap head hanya “menulis” ke potongan dimensinya sendiri. Proyeksi output $W_o$ mencampur hasil semua head sebelum dimasukkan ke residual stream.

Penjelasan

Yang diminta. Causal self-attention dengan n_heads kepala, masing-masing berdimensi hd = D / n_heads, lalu proyeksi output Wo.

Langkah di kode:

Proyeksi. X @ Wq, X @ Wk, X @ Wv, masing-masing (T, D).
Pecah per head. split = lambda M: M.reshape(T, H, hd).transpose(1, 0, 2). reshape(T, H, hd) memotong setiap baris D menjadi H potongan berukuran hd: head h memakai kolom h*hd .. (h+1)*hd. transpose(1, 0, 2) memindahkan H ke depan menjadi (H, T, hd), sehingga H berperan sebagai dimensi batch.
Attention per head sekaligus. Q @ K.transpose(0, 2, 1) / np.sqrt(hd) berbentuk (H, T, T): satu matriks skor per head. Skalanya $\sqrt{hd}$ , bukan $\sqrt D$ , karena dot product-nya hanya sepanjang hd.
Mask kausal + softmax, sama seperti soal sebelumnya. Mask (T, T) di-broadcast ke semua head.
Gabungkan. (w @ V) berbentuk (H, T, hd) → transpose(1, 0, 2) → (T, H, hd) → reshape(T, D). Ini kebalikan persis dari langkah 2.
@ Wo mencampur hasil semua head.

Cek case 1 vs 2. Input dan bobot sama, hanya n_heads berbeda. Dengan 2 head, setiap head hanya melihat 1 dimensi, sehingga pola attention-nya berbeda dan hasilnya berbeda ([0.5, 0.731] vs [0.330, 0.670]). Banyaknya head mengubah apa yang bisa diperhatikan model.

Yang dikembalikan: array (T, D).