Три матриці на голову, чи одна велика матриця

Вигляд на голову

Кожній голові потрібна своя проєкція запиту, проєкція ключа та проєкція значення. Для голови h:

Q_h = X · W_Q^h    де W_Q^h має форму [d_model, head_dim]
K_h = X · W_K^h    де W_K^h має форму [d_model, head_dim]
V_h = X · W_V^h    де W_V^h має форму [d_model, head_dim]

X несе форму входу [batch, seq_len, d_model]. Після проєкції Q_h, K_h, V_h кожне несе форму [batch, seq_len, head_dim].

Об'єднаний вигляд

Матриці для кожної головки розташовані поряд у пам'яті. Єдина об'єднана матриця W_Q форми [d_model, d_model] виробляє всі головки одразу:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

Об'єднане множення матриць виконує один виклик BLAS замість 12. Tensor cores CUDA досягають пікової пропускної здатності на множеннях матриць такого розміру; множення по головах недовикористовували б апаратне забезпечення.

Кількість параметрів

Три злиті матриці W_Q, W_K, W_V, кожна розміром d_model × d_model. Плюс матриця проекції виходу W_O, також d_model × d_model. Для ANDREA-120M:

параметрів на увагу одного шару = 4 × 768² = 2,359,296 ≈ 2.36M
параметрів на 12 шарів      = 12 × 2.36M ≈ 28.3M

Приблизно чверть усіх параметрів ANDREA-120M міститься в проєкціях уваги. Решта трьох чвертей — у підшарі MLP та вбудуваннях.

Назви проєкцій

Дванадцять векторів стають одним

Після обчислень кожного Heads

Кожен head виробляє тензор виводу форми [batch, seq_len, head_dim]. Дванадцять heads виробляють дванадцять таких тензорів. Конкатенація вздовж розмірності ознак з'єднує їх назад разом:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # for ANDREA-120M

Concat повертає назад split. Загальна розмірність ознак повертається до d_model. Немає втрати інформації в розмірностях; різниця полягає в тому, що містить кожен шматок: шматок head 1 відображає вивчений шаблон уваги head 1.

Проекція виходу W_O

Саме конкатенація залишає голови ізольованими: вихід head 4 розташований поруч із виходом head 7, жодна з них не усвідомлює іншу. Проекція виходу W_O форми [d_model, d_model] змішує їх:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

Після W_O кожен вимір вихідних даних несе в собі навчену лінійну комбінацію всіх дванадцяти голів. Інформація вільно тече між головами через цю єдину множення на матрицю.

Чому голови спеціалізуються

Ніщо в архітектурі не змушує голову 4 вивчати час дієслова чи голову 9 — парну пунктуацію. Спеціалізація виникає від тиску градієнтів: під час тренування голови, що роблять повторювані внески, отримують менші градієнти, ніж ті, що роблять унікальні. За тисячі кроків кожна голова осідає в нішу, яка найефективніше зменшує загальні втрати.

Емпірично навчені трансформери демонструють голови, що обробляють: позиційні патерни (голова дивиться на попередній токен), синтаксичні патерни (голова дивиться на відповідну закриваючу дужку), семантичні патерни (голова дивиться на останню згадану іменовану сутність). Жодних міток не потрібно для цієї спеціалізації. Тільки тренувальний сигнал, що поширюється через W_O, сортує голови.

Чому дванадцять голів, а не одна ширша голова

Як CUDA зберігає голівки

Один тензор, переформований

Тренувальний двигун ANDREA microgpt_cuda.cu не виділяє дванадцять окремих буферів для дванадцяти голівок. Він виділяє один злитий тензор & трактує розмірність голівки як шаблон кроку:

// після Q = X · W_Q (один матричний добуток, злитий по голівках)
// Q має форму [batch, seq_len, d_model]
// переформувати до    [batch, seq_len, n_head, head_dim]
// транспонувати до  [batch, n_head, seq_len, head_dim]
// тепер кожна голова є суцільною в двох внутрішніх вимірах

Транспонування переміщує n_head перед seq_len. Чому? Тому що наступна операція (Q_h · K_h^T) потребує, щоб зріз seq_len × head_dim кожної голови був суцільним у пам'яті. CUDA матмули працюють швидше на суцільних тензорах.

Один Ядро, Багато Голов

One Kernel, Many Heads

Один CUDA-ядро уваги виконується паралельно по всіх головах. Кожен блок потоків обробляє одну пару (batch, head); потоки всередині блоку співпрацюють над тайлом seq_len × head_dim. Ядро ніколи не знає, що обробляє кілька голів; паралелізм забезпечує сітка запуску.

Конфігурація Відображає Апаратне Забезпечення

Вибір ANDREA-120M параметрів n_head=12, head_dim=64 узгоджується з тензорними ядрами RTX 4090, які віддають перевагу тайлам matmul, кратним 16. head_dim=64 = 4 × 16 точно відповідає формі тайла. head_dim=32 (ANDREA-12M) також відповідає, але недоутворює тайл. head_dim=72 не відповідала б і змусила використовувати резервні ядра.

Загальна Картина

П’ять кроків. Три матмули торкаються входу (проєкції Q, K, V). Один матмул торкається об’єднаних голів (W_O). Один ядро уваги обробляє всі голови паралельно. ANDREA-120M виконує всі п’ять кроків один раз на шар, дванадцять шарів углиб, кожен прямий прохід.

Читання конвеєра