Три матрицы на голову или одна большая матрица

Вид на голову

Каждой голове нужны свои проекции запроса, ключа и значения. Для головы 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. Тензорные ядра CUDA достигают пиковой пропускной способности на умножениях матриц такого размера; умножения по головам недоиспользовали бы аппаратное обеспечение.

Количество параметров

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

параметров на attention одного слоя = 4 × 768² = 2,359,296 ≈ 2.36M
параметров на 12 слоёв           = 12 × 2.36M ≈ 28.3M

Примерно четверть всех параметров ANDREA-120M содержится в проекциях внимания. Оставшиеся три четверти находятся в подслое MLP и эмбеддингах.

Называние проекций

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

После вычислений каждого головы

Каждый голова производит тензор вывода формы [batch, seq_len, head_dim]. Двенадцать голов производят двенадцать таких тензоров. Конкатенация вдоль размерности признаков объединяет их обратно:

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

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

Проекция выхода W_O

Само по себе объединение оставляет головы изолированными: выход головы 4 находится рядом с выходом головы 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 работают быстрее на контигуозных тензорах.

Один Ядро, Много Голов

Единый CUDA-ядро attention выполняется параллельно по всем головам. Каждый блочок потоков обрабатывает одну пару (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 выполняет все пять шагов один раз на слой, двенадцать слоёв в глубину, при каждом прямом проходе.

Чтение конвейера