tampakwilll/split-multi-head-total

Split Multi-Head Total: Transformer dengan Proyeksi Block-Diagonal per Head

Artikel asli: https://catatanwildanardiansyah.blogspot.com/2026/10/split-multi-head-total-transformer.html

Ringkasan. Split Multi-Head Total (SMHT) adalah varian Transformer decoder yang membatasi semua proyeksi linear di dalam blok, yaitu Q, K, V, W_O, dan dua lapisan FFN, menjadi block-diagonal per head. Vektor berdimensi d_model dipecah menjadi H potongan, dan potongan ke-i hanya diproses oleh bobot milik head ke-i. Dokumen ini menjelaskan strukturnya, menelusuri cara kerjanya dari token masukan sampai logit keluaran langkah demi langkah, lalu menghitung jumlah parameternya.

1. Notasi

Simbol	Arti
B, L	ukuran batch, panjang sekuens (jumlah token)
V	ukuran kosakata
d_model	dimensi vektor tiap token di seluruh model
H	jumlah head; harus membagi d_model dan d_ff
d_head	d_model / H, lebar satu potongan
d_ff, d_ff_head	lebar hidden FFN, dan d_ff / H

Semua proyeksi linear di dalam blok tidak memakai bias.

2. Gagasan inti: dari matriks padat ke block-diagonal

Pada Transformer standar, proyeksi Q, K, dan V dibagi per head, tetapi setiap head membaca seluruh d_model dari inputnya. Setelah attention, matriks W_O berukuran penuh d_model × d_model mencampur semua head. FFN juga berupa dua matriks padat penuh.

SMHT mengganti setiap matriks padat itu dengan matriks block-diagonal berisi H blok kecil. Head ke-i hanya membaca potongan ke-i dari vektor, dan hanya menulis ke potongan ke-i dari keluaran. Dengan H = 4, matriks proyeksi yang efektif berbentuk:

[ A  0  0  0 ]
[ 0  B  0  0 ]
[ 0  0  C  0 ]
[ 0  0  0  D ]

A, B, C, dan D adalah matriks d_head × d_head, masing-masing milik satu head. Angka 0 tidak disimpan sebagai parameter.

3. Struktur model

token_ids                                       (B, L)
  → Embedding token + Embedding posisi           (B, L, d_model)
  → [ Blok SMHT ] × N                            (B, L, d_model)
  → LayerNorm akhir                              (B, L, d_model)
  → LM head: Linear(d_model → V, tanpa bias)     (B, L, V)

Blok SMHT (pre-norm):
  x = x + SplitWO( SplitAttention( LayerNorm(x) ) )
  x = x + SplitFFN( LayerNorm(x) )

Embedding posisi bersifat terpelajar (tabel dengan max_seq baris), bukan sinusoidal. LayerNorm, kedua embedding, dan LM head tetap padat (tidak dipecah). Weight tying antara embedding token dan LM head bersifat opsional.

4. Cara kerja langkah demi langkah

Agar mudah dilacak, bagian ini memakai model mini: V = 100, max_seq = 16, d_model = 8, H = 4 (jadi d_head = 2), d_ff = 16 (jadi d_ff_head = 4), satu blok, B = 1, dan masukan 4 token (L = 4). Pada model sebenarnya hanya angkanya yang lebih besar; langkahnya sama. Bentuk tensor ditulis dalam urutan (B, ...).

Langkah 0: masukan

token_ids berbentuk (1, 4), berisi bilangan bulat dalam rentang [0, V), misalnya [17, 52, 3, 88].

Langkah 1: embedding token → (1, 4, 8)

Tabel E berukuran 100 × 8 dibaca per baris: token 17 menjadi baris ke-17 dari E, dan seterusnya. Setiap token kini berupa vektor 8 dimensi.

Langkah 2: embedding posisi, lalu dijumlahkan → (1, 4, 8)

Tabel P berukuran 16 × 8 dibaca pada baris 0 sampai 3 (posisi tiap token) dan dijumlahkan elemen demi elemen: x0 = E[token] + P[posisi]. Mulai titik ini, 8 dimensi tiap token dibagi menjadi 4 potongan: dimensi [0,1] milik head 0, [2,3] milik head 1, [4,5] milik head 2, dan [6,7] milik head 3. Pembagian ini tetap sama di seluruh blok.

Langkah 3: blok SMHT, sub-lapisan attention

3a. LayerNorm pertama → (1, 4, 8). Untuk tiap token secara terpisah, dihitung rata-rata μ dan varians σ² atas 8 dimensinya, lalu (x − μ) / √(σ² + ε) · γ + β. Parameter γ dan β berukuran 8 dan tidak dipecah per head. Perhatikan bahwa μ dan σ² dihitung dari seluruh dimensi, jadi semua potongan berbagi dua angka ini.

3b. Pecah menjadi head → (1, 4, 4, 2). Vektor 8 dimensi di-reshape menjadi 4 potongan berisi 2 dimensi. Langkah ini hanya mengubah bentuk tensor dan tidak ada perhitungan. Sumbu ketiga sekarang adalah sumbu head.

3c. Proyeksi Q, K, V per head → masing-masing (1, 4, 4, 2). Head h punya tiga matriks 2 × 2 sendiri: Wq[h], Wk[h], Wv[h]. Untuk tiap token:

q_h = x_h · Wq[h]
k_h = x_h · Wk[h]
v_h = x_h · Wv[h]

Semua head dihitung serentak dengan satu einsum atas bobot berbentuk (H, d_head, d_head):

# head h memakai Wq[h]
q = torch.einsum('blhd,hde->bhle', x, Wq)

Head 0 hanya membaca dimensi 0 dan 1, sehingga tidak melihat dimensi 2 sampai 7.

3d. Skor attention dan softmax → (1, 4, 4, 4). Untuk tiap head: S_h = q_h · k_hᵀ / √d_head, yaitu matriks 4 × 4 berisi kemiripan setiap pasang posisi, dibagi √2. Lalu diterapkan mask kausal: elemen pada posisi kolom s yang lebih besar dari baris t diisi −∞, agar token tidak melihat token sesudahnya. Setelah softmax per baris, baris ke-t berisi bobot perhatian atas posisi 0..t yang jumlahnya 1. Tiap head memiliki peta perhatiannya sendiri, dihitung hanya dari potongannya.

3e. Menimbang value → (1, 4, 4, 2). A_h = softmax(S_h) · v_h. Setiap token kini berisi campuran value dari token-token sebelumnya, di dalam potongannya.

3f. Konkatenasi head → (1, 4, 8). Sumbu head dikembalikan ke samping sumbu d_head lalu digabung, sehingga hasil head h kembali menempati dimensi 2h dan 2h+1. Tiap potongan tetap berada di posisinya semula.

3g. W_O per head → (1, 4, 8). Hasil dipecah lagi menjadi (1, 4, 4, 2) dan tiap potongan dikalikan Wo[h] (2 × 2): y_h = A_h · Wo[h], lalu digabung kembali. Tidak ada matriks 8 × 8 yang mencampur potongan sesudahnya.

3h. Koneksi residual → (1, 4, 8). x1 = x0 + y, penjumlahan elemen demi elemen.

Langkah 3 (lanjutan): sub-lapisan FFN

3i. LayerNorm kedua → (1, 4, 8). Sama seperti 3a, dengan γ dan β miliknya sendiri.

3j. Pecah menjadi head → (1, 4, 4, 2). Sama seperti 3b.

3k. FFN per head → (1, 4, 4, 2). Head h punya dua matriks: W1[h] berukuran 2 × 4 dan W2[h] berukuran 4 × 2.

u_h = GELU( x_h · W1[h] )
z_h = u_h · W2[h]

Di sini u_h berbentuk (1, 4, 4), yaitu potongan 2 dimensi dilebarkan menjadi 4, dan z_h berbentuk (1, 4, 2), yaitu dikembalikan menjadi 2 dimensi. Bobot seluruh head berbentuk W1: (4, 2, 4) dan W2: (4, 4, 2).

3l. Konkatenasi dan residual → (1, 4, 8). Hasil z_h digabung menjadi (1, 4, 8), lalu x2 = x1 + z. Bila ada N blok, x2 menjadi masukan blok berikutnya, yang bobotnya terpisah dari blok ini.

Langkah 4: LayerNorm akhir → (1, 4, 8)

Sama seperti 3a, dengan parameter sendiri.

Langkah 5: LM head → (1, 4, 100)

Matriks padat U berukuran 8 × 100 (tanpa bias) dikalikan ke tiap vektor token: logits = x · U. Pada langkah inilah seluruh 8 dimensi, dari semua potongan, digabung menjadi 100 skor. Bila weight tying dipakai, U = Eᵀ.

Langkah 6: membaca keluaran

Logit pada posisi t adalah 100 skor untuk token yang mungkin muncul sesudah posisi t. Softmax atas sumbu terakhir mengubahnya menjadi distribusi probabilitas. Saat menghasilkan teks, distribusi pada posisi terakhir (L − 1) dipakai untuk memilih token berikutnya.

5. Jumlah parameter

• Q, K, V: standar 3 · d_model²; SMHT 3 · d_model · d_head. Model mini: 192 → 48.
• W_O: standar d_model²; SMHT d_model · d_head. Model mini: 64 → 16.
• FFN: standar 2 · d_model · d_ff; SMHT 2 · d_model · d_ff_head. Model mini: 256 → 64.
• Total satu blok: pada SMHT setiap proyeksi bernilai 1/H dari standar. Model mini: 512 → 128.
• LayerNorm (2 per blok + 1 akhir): 2 · d_model per LayerNorm, sama untuk keduanya. Model mini: 3 × 16 = 48.
• Embedding token, embedding posisi, LM head: V · d_model, max_seq · d_model, V · d_model, sama untuk keduanya. Model mini: 800, 128, 800.

Total model mini dengan SMHT: 800 + 128 + 128 + 48 + 800 = 1.904 parameter. Bila H dan d_model besar, bagian blok menjadi sangat kecil dibanding embedding dan LM head.

6. Sifat struktural

6.1 H Transformer sempit yang berjalan paralel

Karena semua bobot block-diagonal, potongan ke-i berevolusi seperti Transformer berlebar d_head dengan satu head attention dan hidden FFN selebar d_ff_head. SMHT dapat dibaca sebagai H model sempit yang berbagi satu lintasan komputasi.

6.2 Titik kopling antar potongan

1. LayerNorm (langkah 3a, 3i, 4): rata-rata dan varians dihitung atas seluruh d_model, sehingga semua potongan berbagi dua angka per token.
2. LM head (langkah 5): karena linear atas hasil konkatenasi, logit kosakata ke-v adalah jumlah kontribusi tiap potongan: logit_v = Σ_i ⟨U_v[i], z_i⟩.

Koneksi residual hanya menjumlahkan dimensi yang sama dengan dimensi yang sama, jadi tidak mencampur potongan. Menambah jumlah blok juga tidak menambah pencampuran, kecuali lewat LayerNorm.

6.3 Kasus H = 1

Seluruh bobot menjadi matriks padat d_model × d_model, dan FFN-nya d_model → d_ff → d_model. Ini identik dengan Transformer standar satu head. Untuk H > 1, SMHT dan Transformer multi-head standar adalah dua arsitektur yang berbeda, sebab pada yang standar tiap head membaca seluruh d_model dan W_O mencampur semua head.

6.4 Kasus H = d_model (d_head = 1)

Setiap potongan berupa satu channel skalar. Q, K, dan V adalah skalar: dengan a = Wq · Wk, skor antara posisi t dan s adalah a · x_t · x_s. FFN-nya adalah fungsi skalar GELU(w1 · x) · w2 bila d_ff_head = 1.

7. Ruang lingkup

Dokumen ini hanya menjelaskan struktur dan cara kerja arsitektur SMHT. Kualitas model, kecepatan, dan perbandingan dengan Transformer standar tidak dibahas dan tidak diklaim.

Referensi

1. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS.
2. Ba, J. L., Kiros, J. R., Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450.
3. Hendrycks, D., Gimpel, K. (2016). Gaussian Error Linear Units (GELUs). arXiv:1606.08415.

Catatan penulisan: naskah ini disusun dengan bantuan AI, dan rumus serta jumlah parameternya dihitung ulang secara manual.

---

Ahmad Wildan Ardiansyah

D4 Teknik Jaringan Telekomunikasi Digital

Teknik Elektro

Politeknik Negeri Malang