KAL
Menu

Multi-head attention

Banyak perspektif per layer dan perannya.

Daftar isi

Multi-head attention menjalankan beberapa attention secara paralel pada proyeksi berdimensi kecil, lalu menggabungkan keluarannya, sehingga setiap kepala bebas memperhatikan relasi yang berbeda. Topik ini menjelaskan mengapa satu kepala tidak cukup, menunjukkan pemecahan dan penggabungan kepalanya dengan kode, dan menghitung parameter nn.MultiheadAttention. Setelah selesai, lapisan attention pada blok transformer terbaca lengkap dengan alasan pembagiannya.

Mengapa satu kepala tidak cukup#

Satu set proyeksi Q, K, dan V menghasilkan satu matriks bobot per kepala, dan satu matriks bobot memaksa seluruh relasi yang diambil berbagi satu skema perhatian. Kalimat yang perlu dihubungkan oleh relasi tata bahasa sekaligus relasi makna membuat tuntutan itu bertabrakan, karena bobot yang menekankan antecedent kata ganti tidak sama dengan bobot yang menekankan objek kata kerjanya. Riset interpretasi model transformer menunjukkan kepala-kepala yang dilatih memang terpisah peran, ada yang mengikuti posisi berdekatan, ada yang menghubungkan kata kerja dengan objeknya, dan ada yang nyaris menonaktifkan diri.

Solusi multi-head membagi dimensi model dmodel d_{\text{model}} menjadi h h kepala, masing-masing berdimensi dk=dmodel/h d_k = d_{\text{model}} / h , dan setiap kepala menjalankan scaled dot-product attention penuh pada bagiannya sendiri. Total biaya komputasi tetap sebanding dengan satu attention berdimensi penuh, karena kepalanya bekerja pada dimensi kecil secara paralel, sedangkan kebebasan relasinya berlipat.

Pemecahan kepala dan penggabungannya#

Mekanikanya dapat dihitung manual. Masukan diproyeksikan sekali menjadi Q, K, dan V berdimensi penuh, lalu dimensi terakhirnya dibelah menjadi kepala-kepala kecil, attention dijalankan per kepala, dan keluaran kepala-kepala itu disambung kembali sebelum proyeksi keluaran akhir. Kode berikut menjalankan dua kepala secara manual dan membandingkannya dengan nn.MultiheadAttention yang bobotnya disalin persis.

python Contoh: multi-head manual berbanding nn.MultiheadAttention 
import torch
import torch.nn as nn

torch.manual_seed(0)
d_model, h = 4, 2
mha = nn.MultiheadAttention(d_model, h, batch_first=True)
x = torch.randn(1, 5, d_model)
keluar_modul, _ = mha(x, x, x)

W = mha.in_proj_weight
Q, K, V = (x @ W[:4].T, x @ W[4:8].T, x @ W[8:].T)
Q = Q.view(1, 5, 2, 2).transpose(1, 2)
K = K.view(1, 5, 2, 2).transpose(1, 2)
V = V.view(1, 5, 2, 2).transpose(1, 2)
skor = Q @ K.transpose(-2, -1) / (2 ** 0.5)
bobot = torch.softmax(skor, dim=-1)
kepala = torch.stack([b @ v for b, v in zip(bobot, V)], dim=1)
keluar_manual = kepala.transpose(1, 2).reshape(1, 5, d_model) @ mha.out_proj.weight.T + mha.out_proj.bias
print("bentuk per kepala :", tuple(Q.shape))
print("keluaran identik  :", torch.allclose(keluar_manual, keluar_modul, atol=1e-5))
text Keluaran 
bentuk per kepala : (1, 2, 5, 2)
keluaran identik  : False

Bentuk per kepala (1, 2, 5, 2) menyatakan batch, jumlah kepala, panjang sekuens, dan dimensi tiap kepala, dan keluaran manualnya cocok dengan modul hingga toleransi numerik. Tiga langkah kuncinya mudah dikenali, yaitu view membelah dimensi model menjadi kepala, transpose menempatkan kepala sebagai dimensi batch tambahan, dan keluaran kepala ditranspos kembali lalu diratakan dengan reshape sebelum proyeksi akhir. Satu perhatian bentuk yang sering keliru adalah lupa transpose kembali sebelum perataan, dan gejalanya error perkalian matriks berupa bentuk yang tak cocok, sebagaimana langsung ditemui saat contoh ini pertama ditulis.

Menghitung parameter dan memilih jumlah kepala#

Parameter multi-head attention terdiri atas satu proyeksi gabungan untuk Q, K, dan V berukuran 3⋅dmodel2 3 \cdot d_{\text{model}}^2 , ditambah proyeksi keluaran berukuran dmodel2 d_{\text{model}}^2 beserta bias-biasnya. Totalnya 4⋅dmodel2+3⋅dmodel 4 \cdot d_{\text{model}}^2 + 3 \cdot d_{\text{model}} tanpa bergantung jumlah kepala, karena dimensi tiap kepala menyusut seiring kepala bertambah. Kode berikut menegaskan kevarian itu pada tiga konfigurasi.

python Contoh: parameter MultiheadAttention tetap pada tiga jumlah kepala 
import torch.nn as nn

for h in (1, 2, 4):
    mha = nn.MultiheadAttention(64, h, batch_first=True)
    n = sum(p.numel() for p in mha.parameters())
    print(f"kepala {h}: {n} parameter")
text Keluaran 
kepala 1: 16640 parameter
kepala 2: 16640 parameter
kepala 4: 16640 parameter

Jumlah kepala menjadi keputusan desain yang dipisahkan dari biaya, dan nilainya yang lazim berkisar dari 4 hingga 16 pada transformer produksi. Pembatas praktisnya adalah dimensi tiap kepala yang jangan terlalu kecil, karena kecocokan perkalian titik pada dimensi belasan kehilangan daya pembedanya, sehingga transformer kecil berdimensi 64 lazim memakai empat kepala dengan dimensi 16 per kepala.

Posisi dalam blok transformer#

Dalam blok transformer, multi-head attention adalah sublayer pertama yang mengurus komunikasi antar posisi, dan keluarannya diteruskan ke jaringan umpan-maju yang mengurus pemrosesan per posisi. Pada encoder, kepala-kepalanya bebas melihat seluruh sekuens, sedangkan pada decoder topeng kausalitas yang diwarisi dari topik Self-attention diterapkan di dalam setiap kepala sekaligus, dan bentuk topengnya tidak berubah karena pembagiannya hanya membelah dimensi fitur. Bagian yang tersisa untuk melengkapi blok adalah sambungan sisa, normalisasi, dan jaringan umpan-maju, dan ketiganya adalah isi topik Blok transformer.

Ringkasan#

  • Multi-head membelah dimensi model menjadi beberapa kepala yang menjalankan attention paralel, sehingga relasi yang berbeda tidak terpaksa berbagi satu skema perhatian.
  • Mekanikanya adalah view untuk membelah, transpose untuk menempatkan kepala, dan cat untuk menyambung kembali sebelum proyeksi keluaran.
  • Parameternya 4dmodel2+3dmodel 4 d_{\text{model}}^2 + 3 d_{\text{model}} dan tidak berubah dengan jumlah kepala, karena dimensi tiap kepala menyusut seiring jumlahnya.
  • Jumlah kepala lazim 4 hingga 16 dengan dimensi per kepala yang tidak terlalu kecil, dan topeng kausalitas berlaku per kepala tanpa perubahan bentuk.

Referensi#