Mechanismus pozornosti: jak se model naučí, kam se dívat

Nechť d, m, n ∈ ℕ, d, m, n ≥ 1. Nechť q ∈ ℝ^d je dotaz, k₁, …, k_n ∈ ℝ^d jsou klíče a v₁, …, v_n ∈ ℝ^m jsou hodnoty (klíč k_i a hodnota v_i tvoří i-tou položku). Pozornost (škálovaná, se skalárním součinem) přiřadí dotazu vektor

Attention(q; k₁, …, k_n; v₁, …, v_n) = ∑_i=1ⁿ a_i v_i ∈ ℝ^m,

kde pro i = 1, …, n:

1. skóre s_i = q · k_i/√d, přičemž q · k_i = ∑_r=1^d q_r(k_i)_r je skalární součin;
2. váhy (a₁, …, a_n) = softmax(s₁, …, s_n), tj. a_i = exp(s_i)/∑_j=1ⁿ exp(s_j) (funkce softmax je zavedená v dílu Vícetřídní klasifikace a softmax).

Platí a_i > 0 a ∑_i=1ⁿ a_i = 1, takže výsledek pozornosti je vážený průměr hodnot, ve kterém má větší váhu hodnota té položky, jejíž klíč má s dotazem větší skalární součin.

V neuronové síti vznikají dotaz, klíče a hodnoty z reprezentací vstupu: nechť x ∈ ℝ^p je reprezentace toho, kdo se ptá, x₁, …, x_n ∈ ℝ^p′ reprezentace položek a W_Q ∈ ℝ^d×p, W_K ∈ ℝ^d×p′, W_V ∈ ℝ^m×p′ naučené matice. Pak q = W_Qx, k_i = W_Kx_i a v_i = W_Vx_i pro i = 1, …, n.

Dimenze d = 2, dotaz q = (2; 1), klíče k₁ = (1; 0), k₂ = (0; 1), k₃ = (−1; 0) a hodnoty (pro jednoduchost čísla) v = (10; 20; 30).

Barvy odpovídají členům vzorce výše.

Dotaz se nejvíc podobá prvnímu klíči, takže výsledek je blízko jeho hodnotě 10, ale je do něj přimíchaná i dvacítka. Obyčejný slovník by vrátil přesně 10; čím víc se skóre prvního klíče od ostatních odliší, tím víc se mu měkký slovník přiblíží.

Nechť 𝒫 = {a, b, …, h} je množina osmi písmen a 𝒞 = {0, 1, …, 9} množina deseti číslic. Pro písmeno π ∈ 𝒫 označme e_π ∈ {0, 1}⁸ jeho vektor one-hot (jednička na pozici písmena, jinde nuly; viz díl Embeddingy a word2vec) a pro číslici c ∈ 𝒞 obdobně u_c ∈ {0, 1}¹⁰.

Příklad úlohy je šestice dvojic (π_i, c_i) ∈ 𝒫 × 𝒞, i = 1, …, 6, s navzájem různými písmeny (π_i ≠ π_j pro i ≠ j), spolu s dotazovaným písmenem π* = π_i* pro nějaké i* ∈ {1, …, 6}. Správná odpověď je c* = c_i*. Vstupy modelu jsou

x_i = (e_πi, u_ci) ∈ {0, 1}¹⁸   (i = 1, …, 6),   x_dotaz = e_π* ∈ {0, 1}⁸,

kde (e, u) značí spojení dvou vektorů za sebe. Model s pozorností s dimenzí d = 16 má parametry W_Q ∈ ℝ^d×8, W_K, W_V ∈ ℝ^d×18, W_O ∈ ℝ^10×d a b_O ∈ ℝ¹⁰ a počítá:

1. q = W_Qx_dotaz,   k_i = W_Kx_i,   v_i = W_Vx_i pro i = 1, …, 6;
2. o = Attention(q; k₁, …, k₆; v₁, …, v₆) = ∑_i=1⁶ a_iv_i ∈ ℝ^d (definice pozornosti výše);
3. y = softmax(W_Oo + b_O) ∈ ℝ¹⁰; složka y_c je pravděpodobnost, kterou model dává číslici c ∈ 𝒞.

Ztráta na jednom příkladu je

L = −ln y_c*,

tj. křížová entropie (s přirozeným logaritmem, v jednotkách nat) mezi správnou odpovědí ve tvaru one-hot u_c* a předpovědí y (křížová entropie je zavedená v dílu Teorie informace). Průměrování je týž model, jen krok 2 počítá s pevnými vahami a_i = 1/6 pro všechna i.

Model, který nic neví a dá všem deseti číslicím y_c = 1/10, má ztrátu −ln 1/10 = ln 10 ≈ 2,303 — to je čárkovaná čára v obrázku níže.

Matice pozornosti mají stejný tvar pro všech šest dvojic — jedno W_K slouží každé pozici. Plně propojená síť má pro každou pozici vlastní kus matice W₁, a proto má parametrů 9krát víc.

Zjednodušeně předpokládejme, že ostatních pět dvojic má stejné skóre a správná je o Δ výš. Pak

Správná dvojice tedy dostane asi 134krát větší váhu než každá jiná. K tomu stačí, aby skóre (už vydělené √d) bylo o necelých pět jednotek vyšší — jinými slovy, aby model měl matice W_Q, W_K naladěné tak, že shodné písmeno dává výrazně větší skalární součin než neshodné.

Nechť d ∈ ℕ, d ≥ 1, a nechť q = (q₁, …, q_d) a k = (k₁, …, k_d) jsou náhodné vektory v ℝ^d, jejichž složky q₁, …, q_d, k₁, …, k_d jsou navzájem nezávislé náhodné veličiny a pro každé j = 1, …, d platí

E[q_j] = E[k_j] = 0,   Var(q_j) = Var(k_j) = 1.

Pak E[q · k] = 0 a

Var(q · k) = d,   a proto   Var(q · k/√d) = 1.

Důkaz. Pro jeden člen: z nezávislosti E[q_jk_j] = E[q_j] E[k_j] = 0 a E[q_j²] = Var(q_j) + (E[q_j])² = 1, stejně E[k_j²] = 1. Proto

Členy q_jk_j pro různá j jsou nezávislé (každý závisí na jiné dvojici složek), takže E[q · k] = ∑_j E[q_jk_j] = 0 a rozptyly se sčítají: Var(q · k) = ∑_j=1^d 1 = d. Pro konstantu c je Var(cX) = c² Var(X), takže dělení √d dá rozptyl d/d = 1. ∎

Derivace softmaxu: ∂a_i/∂s_j = a_i([i = j] − a_j), kde [i = j] je 1 pro i = j a 0 jinak; tedy pro váhu podle vlastního skóre a_i(1 − a_i). Dosadíme průměrné největší váhy z tabulky pro d = 256:

Gradient, který přes softmax proteče ke skóre (a dál k W_Q, W_K), se těmito činiteli násobí. Čím jistější softmax, tím menší krok učení — u váhy blízko 1 se skoro zastaví.

Nechť n, d, d_k, d_v ∈ ℕ jsou kladná čísla a x₁, …, x_n ∈ ℝ^d vektory tokenů jedné posloupnosti; uložme je jako řádky matice X ∈ ℝ^n×d. Nechť W_Q, W_K ∈ ℝ^d×d_k a W_V ∈ ℝ^d×d_v jsou naučené matice. Sebepozornost (self-attention) je zobrazení ℝ^n×d → ℝ^n×d_v, které z téže posloupnosti spočte pro každý token j = 1, …, n dotaz, klíč i hodnotu

q_j = W_Q^Tx_j,   k_j = W_K^Tx_j,   v_j = W_V^Tx_j

(tj. řádky matic Q = XW_Q, K = XW_K, V = XW_V) a nový vektor tokenu j určí pozorností jeho dotazu na všechny tokeny:

z_j = Attention(q_j; k₁, …, k_n; v₁, …, v_n) = ∑_i=1ⁿ a_ji v_i,   (a_j1, …, a_jn) = softmax(q_j · k₁/√d_k, …, q_j · k_n/√d_k).

Nechť dále h ∈ ℕ, h ≥ 1, a pro každé r = 1, …, h má hlava r vlastní matice W_Q^(r), W_K^(r), W_V^(r) uvedených rozměrů; head_r ∈ ℝ^n×d_v je výsledek sebepozornosti s nimi. Nechť W_O ∈ ℝ^hd_v×d. Vícehlavá pozornost je

MultiHead(X) = [head₁, …, head_h] W_O ∈ ℝ^n×d,

kde [·, …, ·] značí matice položené vedle sebe (spojení po sloupcích, výsledek má rozměr n × hd_v). Každý token tedy dostane nový vektor, který je v každé hlavě váženým průměrem hodnot všech tokenů posloupnosti.

Původní transformer (Vaswani a kol., 2017) měl vektory tokenů dimenze d = 512 a h = 8 hlav, tedy d_k = 512/8 = 64. Pro n = 1 000 tokenů a jednu hlavu:

Počet parametrů na délce textu nezávisí — stejné matice slouží každé pozici —, počet operací roste s n².

Nechť originál má n slov a kodér pro i-té z nich spočetl stav h_i ∈ ℝ^d_h, i = 1, …, n. Nechť dekodér právě píše t-té slovo překladu (t ≥ 1) a jeho předchozí stav je s_t−1 ∈ ℝ^d_s. Nechť W ∈ ℝ^p×d_s, U ∈ ℝ^p×d_h a w ∈ ℝ^p jsou naučené parametry (p je počet neuronů skryté vrstvy) a tanh se aplikuje po složkách. Aditivní pozornost spočte pro i = 1, …, n:

1. skóre malou sítí s jednou skrytou vrstvou: e_ti = w^T tanh(W s_t−1 + U h_i) ∈ ℝ;
2. váhy (a_t1, …, a_tn) = softmax(e_t1, …, e_tn);
3. kontextový vektor, který dekodér použije při volbě t-tého slova: c_t = ∑_i=1ⁿ a_ti h_i ∈ ℝ^d_h.

V řeči tohoto dílu je s_t−1 dotaz a h_i klíč i hodnota zároveň. Transformer nahradil síť pro skóre skalárním součinem q · k_i/√d: dá podobné výsledky, ale všechna skóre se spočtou jedním násobením matic, což je na grafické kartě mnohem rychlejší.