Scaled dot-product attention

Implementasikan inti transformer:

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V

Input bisa punya dimensi batch di depan. Softmax dilakukan di sepanjang axis key (axis terakhir dari skor), dan harus stabil walau skornya besar.

Q = [[1, 0]], K = [[1, 0], [0, 1]], V = [[10, 0], [0, 10]]
skor = [0.707, 0], bobot = [0.67, 0.33]  →  [[6.70, 3.30]]

causal mask bagi √d

query ↓ / key →	kucing	itu	tidur	di	atas	kasur
kucing	0.20	0.16	0.17	0.15	0.15	0.16
itu	0.12	0.12	0.12	0.16	0.21	0.27
tidur	0.35	0.14	0.17	0.11	0.11	0.12
di	0.11	0.11	0.11	0.16	0.22	0.29
atas	0.12	0.13	0.13	0.17	0.20	0.25
kasur	0.17	0.15	0.15	0.16	0.17	0.19

Baris “tidur”: paling memperhatikan kucing (35%). Outputnya adalah rata-rata berbobot dari vektor V semua token dengan bobot di baris ini.

Attention(Q, K, V) = softmax(QKᵀ / √d) V. Setiap baris dijumlahkan menjadi 1. Matikan "bagi √d" dan lihat baris menjadi lebih "tajam". Nyalakan causal mask: token hanya boleh melihat dirinya dan token sebelumnya (dipakai GPT).

Query, key, value

Bayangkan setiap token bertanya (query): “informasi apa yang saya butuhkan?”. Setiap token juga punya label (key): “saya berisi informasi tentang apa”. Skor $q_i\cdot k_j$ mengukur seberapa cocok pertanyaan token $i$ dengan label token $j$ . Softmax mengubah skor menjadi bobot, lalu output token $i$ adalah rata-rata berbobot dari isi (value) semua token.

Kenapa dibagi √d

Jika komponen $q$ dan $k$ independen dengan varians 1, maka $q\cdot k$ punya varians $d$ . Untuk $d = 64$ , skornya berkisar $\pm 8$ atau lebih. Softmax dari skor sebesar itu hampir one-hot, dan gradiennya nyaris nol. Membagi dengan $\sqrt{d}$ mengembalikan varians skor ke 1. Matikan opsi “bagi √d” di widget untuk melihat efeknya.

Bentuk

(T, d) @ (d, T) → (T, T): setiap pasangan token punya satu skor. Karena itu biaya attention kuadratik terhadap panjang konteks, dan itulah alasan konteks panjang mahal.

Penjelasan

Yang diminta. $\text{softmax}(QK^\top/\sqrt d)\,V$ , untuk input yang boleh punya dimensi batch di depan.

Tiga langkah, masing-masing dengan maknanya:

Skor kecocokan. scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(d). Baris i kolom j adalah $q_i\cdot k_j / \sqrt d$ : seberapa cocok “pertanyaan” token i dengan “label” token j. Bentuknya (..., Tq, d) @ (..., d, Tk) → (..., Tq, Tk). np.swapaxes(K, -1, -2) menukar dua axis terakhir dan bekerja untuk berapa pun dimensi batch; K.T hanya benar untuk 2D.
Bobot. softmax(scores) di axis terakhir (axis key), sehingga setiap baris berjumlah 1 dan menjawab “token i membagi perhatiannya ke token mana saja”.
Campuran. @ V: output token i adalah rata-rata berbobot dari semua baris V, dengan bobot dari baris i.

Kenapa dibagi √d. Dot product dari vektor berdimensi d dengan komponen acak bervarians 1 punya varians d. Tanpa pembagi, skornya besar dan softmax menjadi hampir one-hot, sehingga gradiennya hampir hilang.

Cek case 1. Skor $= [1/\sqrt2,\ 0] = [0.707, 0]$ , softmax $= [0.670, 0.330]$ , output $= 0.670\cdot[10,0] + 0.330\cdot[0,10] = [6.70, 3.30]$ .

Case 2 (Q = 0). Semua skor 0, jadi bobotnya rata dan output adalah rata-rata biasa dari V: $(1 + 3)/2 = 2$ .

Yang dikembalikan: array (..., Tq, dv).