Tiga Matriks Per Head, Atau Satu Matriks Besar

Pandangan Per-Head

Setiap head membutuhkan proyeksi query, proyeksi key, & proyeksi value sendiri. Untuk head h:

Q_h = X · W_Q^h    di mana W_Q^h memiliki bentuk [d_model, head_dim]
K_h = X · W_K^h    di mana W_K^h memiliki bentuk [d_model, head_dim]
V_h = X · W_V^h    di mana W_V^h memiliki bentuk [d_model, head_dim]

X membawa bentuk masukan [batch, seq_len, d_model]. Setelah proyeksi, Q_h, K_h, V_h masing-masing membawa bentuk [batch, seq_len, head_dim].

Pandangan Terpadu

Matriks per-head berada berdampingan di memori. Matriks terpadu tunggal W_Q dengan bentuk [d_model, d_model] menghasilkan semua head sekaligus:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

Matmul terpadu mengirimkan satu panggilan BLAS daripada 12. CUDA tensor cores mencapai throughput puncak pada matmul berukuran ini; matmul per-head akan kurang memanfaatkan perangkat keras.

Jumlah Parameter

Tiga matriks yang digabung W_Q, W_K, W_V, masing-masing d_model × d_model. Ditambah proyeksi output W_O, juga d_model × d_model. Untuk ANDREA-120M:

params per lapisan attention = 4 × 768² = 2,359,296 ≈ 2.36M
params di seluruh 12 lapisan      = 12 × 2.36M ≈ 28.3M

Sekitar seperempat dari total parameter ANDREA-120M berada di proyeksi atensi. Tiga perempat sisanya berada di sublayer MLP & embeddings.

Menamai Proyeksi

Dua Belas Vektor Menjadi Satu

Setelah Setiap Head Menghitung

Setiap head menghasilkan tensor keluaran dengan bentuk [batch, seq_len, head_dim]. Dua belas head menghasilkan dua belas tensor seperti itu. Penggabungan sepanjang dimensi fitur menyusunnya kembali bersama:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # untuk ANDREA-120M

Concat membalikkan split. Dimensi fitur total kembali ke d_model. Tidak ada kehilangan informasi dalam dimensi; perbedaannya terletak pada apa yang terkandung dalam setiap chunk: chunk head 1 mencerminkan pola perhatian yang dipelajari head 1.

Proyeksi Output W_O

Konkatenasi saja membuat head-head tetap terisolasi: output head 4 berada di sebelah output head 7, keduanya tidak menyadari yang lain. Proyeksi output W_O dengan bentuk [d_model, d_model] mencampurnya:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

Setelah W_O, setiap dimensi output membawa kombinasi linear yang dipelajari dari semua dua belas head. Informasi mengalir bebas antar head melalui perkalian matriks tunggal ini.

Mengapa Head Mengkhususkan Diri

Tidak ada yang dalam arsitektur memaksa head 4 untuk mempelajari bentuk kata kerja atau head 9 untuk mempelajari tanda baca berpasangan. Spesialisasi muncul dari tekanan gradien: selama pelatihan, head yang berkontribusi secara redundan menerima gradien yang lebih kecil daripada head yang berkontribusi secara unik. Selama ribuan langkah, setiap head menetap ke dalam ceruk yang paling efektif mengurangi total loss.

Secara empiris, transformer yang dilatih menunjukkan head yang menangani: pola posisional (head melihat token sebelumnya), pola sintaksis (head melihat tanda kurung penutup yang cocok), pola semantik (head melihat entitas bernama terbaru). Tidak ada label yang melatih spesialisasi ini. Sinyal pelatihan saja, yang dipropagasi melalui W_O, menyortir head-head tersebut.

Mengapa Dua Belas Head, Bukan Satu Head yang Lebih Lebar

Bagaimana CUDA Menyimpan Head-Head

Satu Tensor Tunggal, Dibentuk Ulang

Mesin pelatihan ANDREA microgpt_cuda.cu tidak mengalokasikan dua belas buffer terpisah untuk dua belas head. Ia mengalokasikan satu tensor terpadu & memperlakukan dimensi head sebagai pola stride:

// setelah Q = X · W_Q (satu matmul, terpadu di seluruh head)
// Q memiliki bentuk [batch, seq_len, d_model]
// ubah bentuk menjadi    [batch, seq_len, n_head, head_dim]
// transpos menjadi  [batch, n_head, seq_len, head_dim]
// setiap head sekarang berurutan di dua dimensi dalam

Transpos memindahkan n_head ke depan seq_len. Mengapa? Karena operasi berikutnya (Q_h · K_h^T) membutuhkan irisan seq_len × head_dim setiap head berurutan di memori. Matmul CUDA berjalan lebih cepat pada tensor berurutan.

Satu Kernel, Banyak Head

Sebuah kernel CUDA attention tunggal berjalan di seluruh head secara paralel. Setiap thread block menangani satu pasangan (batch, head); thread di dalam block bekerja sama pada tile seq_len × head_dim. Kernel tidak pernah tahu bahwa ia memproses multiple head; grid peluncuran yang menangani paralelisme.

Konfigurasi Mencerminkan Perangkat Keras

Pilihan ANDREA-120M untuk n_head=12, head_dim=64 selaras dengan tensor core RTX 4090, yang lebih menyukai tile matmul dalam kelipatan 16. head_dim=64 = 4 × 16 cocok persis dengan bentuk tile. head_dim=32 (ANDREA-12M) juga cocok tetapi kurang memanfaatkan tile. head_dim=72 tidak akan cocok & akan memaksa kernel fallback.

Gambaran Akhir

Lima langkah. Tiga matmul menyentuh input (proyeksi Q, K, V). Satu matmul menyentuh kepala yang digabungkan (W_O). Satu kernel perhatian menangani semua kepala secara paralel. ANDREA-120M menjalankan kelima langkah sekali per lapisan, dua belas lapisan dalam, setiap forward pass.

Membaca Pipeline