Математика уваги крок за кроком: від вхідної матриці X до softmax і V

Щоразу, коли мовна модель видає токен, знову й знову виконується одне обчислення — увага. Це не фігура мови. На кожному прямому проході, для кожного токена, на кожному шарі виконується та сама арифметика.

Більшість пояснень зупиняються на інтуїції «запитів, ключів і значень», і сама математика лишається чорною скринькою. Ця стаття її відкриває. У ній розібрано, чим насправді є вхідна матриця, звідки беруться Q, K і V, числовий приклад, досить малий, щоб прогнати його самому, і чому це називається self-attention. Після цього ви зможете пройти одну голову уваги вручну.

Яку задачу розв'язує увага

Візьмімо речення: «The animal didn't cross the street because it was tired» («Тварина не перейшла вулицю, бо вона втомилася»).

Дійшовши до токена «it», модель має зрозуміти, до чого стосується це слово. До тварини? До вулиці? Саме собою слово нічого не каже, тож відповідь має прийти з навколишньої послідовності.

Увага розв'язує це, дозволяючи кожному токену дати кожному іншому токену вагу — число, що показує, наскільки йому варто дивитися на цю позицію, поки він будує своє оновлене представлення. Для «it» ваги могли б вийти такими:

"it" → "animal"  : 0.65
"it" → "street"  : 0.20
"it" → "tired"   : 0.10
"it" → others    : 0.05

Ці числа ніхто не пише вручну. Вони виходять з навченого обчислення, побудованого на трьох векторах — Query, Key і Value, а вся формула вміщається в один рядок:

Attention(Q, K, V) = softmax( QKᵀ / √dₖ ) · V

Усе подальше — просто розпакування цього рядка.

Крок 0: вхідна матриця X

Перш ніж з'являться Q, K чи V, потрібна вхідна матриця X.

Кожен токен у «The cat sat» відображається в навчений вектор ембедінгу. За розміру ембедінгу d_model = 4 це виглядає так:

The → [0.2, 0.5, 0.1, 0.7]
cat → [0.8, 0.3, 0.6, 0.2]
sat → [0.4, 0.9, 0.5, 0.1]

Складіть вектори рядок за рядком, і вийде X форми [n_tokens, d_model].

Одну деталь легко пропустити: X не стала від шару до шару. Кожен шар уваги бере вихід попереднього шару як свою X. На першому шарі X — це просто сирий токен плюс інформація про позицію. Вище стеком, скажімо на 12-му шарі GPT-3, що має 96 шарів, X уже кодує глибокий контекстний зміст, накопичений за багато раундів уваги.

Крок 1: обчислюємо Q, K і V

Запити, ключі й значення ніхто не призначає вручну. Усі три обчислюються з X через три окремі навчені матриці ваг:

Q = X · WQ
K = X · WK
V = X · WV

WQ, WK і WV — матриці параметрів, які заморожуються після завершення навчання. Форми сходяться так:

X.shape  = [n, d_model]
WQ.shape = [d_model, d_k]   → Q.shape = [n, d_k]
WK.shape = [d_model, d_k]   → K.shape = [n, d_k]
WV.shape = [d_model, d_v]   → V.shape = [n, d_v]

Різниця важлива, тому ось вона в таблиці:

Символ	Що це	Залежить від входу?
WQ, WK, WV	Навчені матриці проєкції	Ні, фіксовані після навчання
Q, K, V	Обчислювані активації	Так, ідуть за поточною X

Ось чому «bank» у «river bank» (берег річки) і «bank» у «deposit at the bank» (вклад у банку) отримують зовсім різні патерни уваги. Навколишній контекст дає іншу X, яка проходить крізь ті самі матриці ваг і виходить іншими векторами Q, K і V.

Зручна уявна модель — пошуковик:

• Q — те, що шукає токен.
• K — те, що кожен токен повідомляє про себе.
• V — інформація, що повертається в разі збігу.

Q і K вирішують, куди дивитися. V вирішує, що справді йде на вихід.

Крок 2: скалярний добуток QKᵀ

Візьміть Q і помножте на транспоновану K. Вийде квадратна матриця [n, n] сирих оцінок, де елемент у рядку i, стовпці j показує, наскільки сильно токен i звертає увагу на токен j. До підстановки чисел вона виглядає так:

Keys
       cat   eats  fish
     ┌──────┬──────┬──────┐
Q cat│  ?   │  ?   │  ?   │
  eat│  ?   │  ?   │  ?   │
 fish│  ?   │  ?   │  ?   │
     └──────┴──────┴──────┘

Що більший скалярний добуток, то релевантніші токени один одному.

Крок 3: масштабування на √dₖ

Коли d_k велике, сирі скалярні добутки роздуваються. За d_k = 512 оцінки на кшталт [180, 210, 95] — звична річ. Подайте їх у softmax, і на виході буде [0.00, 1.00, 0.00] — цілком сколапсований розподіл. Тоді градієнти згасають, і навчання стає.

Ділення на √d_k тримає оцінки в розумному діапазоні. Саме за це ділення механізм і дістав назву scaled dot-product attention.

Крок 4: softmax перетворює оцінки на ваги

Softmax застосовується рядок за рядком, і кожен рядок стає розподілом імовірностей за всіма токенами. Ось повний розрахунок для оцінок [3.2, 4.0, 1.0, 6.2, 8.0]:

exp(3.2) =     24.53
exp(4.0) =     54.60
exp(1.0) =      2.72
exp(6.2) =    492.75
exp(8.0) =  2,980.96
         ──────────
Σ        =  3,555.56

Token 1:   24.53 / 3555.56 →  0.69%
Token 2:   54.60 / 3555.56 →  1.54%
Token 3:    2.72 / 3555.56 →  0.08%
Token 4:  492.75 / 3555.56 → 13.86%
Token 5: 2980.96 / 3555.56 → 83.83%  ← winner

Токен 5 забирає 83,83%, хоча випереджає токен 4 лише на 1,8 пункту. Саме таке нелінійне підсилення дає моделі змогу впевнено спиратися на найрелевантніші токени.

Порада щодо стабільності: спершу відніміть максимум рядка, як у np.exp(x - x.max()), і лише потім діліть. 64-бітне число з рухомою комою не вміщує exp(800), а після зсуву всі експоненти лишаються в безпечних межах, і підсумкові ймовірності математично ідентичні.

Крок 5: зважена сума з V

Останній крок множить ваги на V і все додає. Для «it» це означає:

output("it") =
    0.65 × V("animal")
  + 0.20 × V("street")
  + 0.10 × V("tired")
  + 0.05 × V(rest)

На виході — новий вектор для «it», у якому здебільшого інформація з «animal», трохи з «street» і зовсім мало з решти. Це контекстуалізоване представлення й іде в наступний шар.

Повний приклад на NumPy

Тут речення «cat eats fish» з d_model = 2:

import numpy as np

X  = np.array([[1, 2], [3, 1], [2, 4]], dtype=float)
WQ = np.array([[1, 0], [0, 1]], dtype=float)  # identity
WK = np.array([[1, 0], [0, 1]], dtype=float)  # identity
WV = np.array([[0, 1], [1, 0]], dtype=float)  # swaps dimensions

Q, K, V = X @ WQ, X @ WK, X @ WV

# Raw scores for "fish" (Q=[2,4]) against each key:
# fish→cat:  [2,4]·[1,2] = 10
# fish→eats: [2,4]·[3,1] = 10
# fish→fish: [2,4]·[2,4] = 20  ← self-attention dominates

scores = Q @ K.T / np.sqrt(2)
e = np.exp(scores - scores.max(-1, keepdims=True))
weights = e / e.sum(-1, keepdims=True)  # softmax row-wise

output = weights @ V
# output[fish] ≈ [3.997, 2.001] ≈ V(fish)
# fish self-attends at ~99.8% in this toy example

У справжній навченій моделі WQ і WK порушили б цю симетрію й дали б змістовну увагу між токенами. Одиничні матриці потрібні лише для того, щоб арифметика лишалася читабельною.

Чому це називається self-attention

Q, K і V беруться з однієї й тієї самої вихідної матриці X. Кожен токен породжує запит, ключ і значення з власного представлення, а потім кожен токен звертає увагу на кожен інший у цій одній послідовності.

Це відрізняє self-attention, як у GPT і LLaMA, від cross-attention, як у моделях encoder-decoder на кшталт T5, де Q береться з однієї послідовності, а K і V — з іншої.

Практичні інженерні нотатки

KV-кеш. Перераховувати K і V для всіх попередніх токенів на кожному кроці генерації коштувало б O(n²). KV-кеш зберігає всі попередні вектори K і V і перевикористовує їх, знижуючи вартість кроку до O(n) в обмін на пам'ять, що зростає з довжиною послідовності.

Розмірність голови. Звична конфігурація — d_k = d_model / num_heads. У GPT-2 це 768 / 12 = 64, у GPT-3 — 12288 / 96 = 128.

Grouped Query Attention (GQA). LLaMA-2, Mistral і Gemma дозволяють кільком головам Query ділити одну голову KV. Це різко скорочує пам'ять під KV-кеш і значною мірою пояснює, чому ці моделі дають раду довгому контексту на споживчому залізі.

Увага в одному реченні

Жодної магії тут немає. Увага — це навчена, диференційовна таблиця пошуку з м'якими вагами, яка наново перераховується для кожного входу й дає представлення, що несуть усю реляційну структуру тексту.