Многоголовость и маски

Головы — это блочная структура

Многоголовое внимание не добавляет вычислений. Оно делит существующую ширину на части и считает внимание в каждой независимо.

$$\text{MHA}(X) = \htmlData{k=concat}{\text{concat}}\big(\htmlData{k=split}{\text{head}_1, \dots, \text{head}_h}\big) W^O, \qquad d = D/h$$

Из прошлого урока: стоимость равна $2n^2D + 4nD^2$ и не содержит $h$. Двадцать четыре головы по $32$ и двенадцать по $64$ стоят одинаково — меняется только то, как ширина поделена.

В нотации индексов разбиение — это переформовка одной оси в две:

$\texttt{'b n (h d) -> b h n d'}$

Именно поэтому einops удобнее, чем view и permute: здесь видно, что $h$ и $d$ — это разложение одной оси, а не две независимые. Порядок множителей в скобках важен, и перепутать (h d) с (d h) — распространённая ошибка, которая не падает.

Зачем несколько голов

Если стоимость та же, что даёт разбиение? Одна голова может внимать только «в одном направлении»: её распределение по ключам одно. Несколько голов дают несколько распределений одновременно.

Ограничение, из которого это следует: у одной головы результат — выпуклая комбинация значений, то есть точка внутри их выпуклой оболочки, с одним набором весов. Чтобы собрать информацию из двух разных мест по разным правилам, нужны два набора весов.

Цена — размер головы. При $D = 768$:

$h$	$d$	что происходит
$1$	$768$	одно распределение, максимум выразительности внутри него
$12$	$64$	стандартный компромисс
$96$	$8$	оценки почти вырождены: $q\cdot k$ по восьми компонентам очень шумно

Последняя строка — реальное ограничение. При очень малом $d$ скалярное произведение вычисляется по нескольким числам, и его дисперсия относительно велика: внимание становится шумным. Отсюда и практический диапазон $d \in [32, 128]$ почти во всех моделях.

Причинная маска

Для авторегрессионной модели токен не должен видеть будущее. Реализуется это прибавлением $-\infty$ к запрещённым оценкам до softmax:

$s_{qk} \leftarrow \begin{cases} s_{qk}, & k \le q \\ -\infty, & k > q \end{cases}$

Почему именно $-\infty$, а не удаление или ноль:

• ноль не подходит: нулевая оценка даёт вес $e^0 = 1$, то есть не запрет, а средний вес. Маскировать нужно до softmax, а не после;
• $-\infty$ даёт ровно нуль после softmax: $e^{-\infty} = 0$, и знаменатель тоже не включает это слагаемое. Нормировка остаётся корректной автоматически;
• удаление не подходит технически: тензор должен остаться прямоугольным, чтобы операции оставались батчевыми.

На практике вместо $-\infty$ берут большое отрицательное число (например $-10^9$ или минимальное значение типа), потому что $-\infty$ в арифметике может дать nan при умножении на нуль. Но с точки зрения математики это именно $-\infty$.

Включите маску в виджете: верхний треугольник заштрихован, и первая строка имеет единственный ненулевой вес — первый токен видит только себя.

делить на √d причинная маска

ключи

запросы

	<s>	кот	сел	на	коврик	.
<s>
кот
сел
на
коврик
.

размер головы d

64

температура T

1

разброс оценок

4.38

энтропия строки

1.209 бит

максимальный вес

1

масштаб градиента

2.47e-1

Обратите внимание на первую строку: у неё энтропия равна нулю независимо от оценок. Первый токен имеет ровно один допустимый ключ — себя, — и после нормировки его вес равен единице принудительно. Никакой информации в этой строке нет, и это не дефект: так устроена авторегрессия.

Сколько работы экономит маска

Наивная реализация считает всю матрицу, а потом половину выбрасывает. Экономии нет вовсе — только трата.

$n$	используется	доля
$4$	$10$ из $16$	$62.5\%$
$128$	$8256$ из $16384$	$50.4\%$
$1024$	$524\,800$ из $1\,048\,576$	$50.05\%$

Доля равна $\frac{n+1}{2n}$ и стремится к половине сверху. То есть при честной реализации причинное внимание должно быть примерно вдвое дешевле полного — и FlashAttention с блочными ядрами эту экономию действительно получает, просто не вычисляя запрещённые блоки.

Другие маски

маска	что запрещает	где
причинная	будущее	декодеры, языковые модели
padding	заполнитель в коротких последовательностях	батчи переменной длины
скользящее окно	всё дальше $w$ токенов	Longformer, Mistral
блочно-разреженная	всё вне выбранных блоков	длинный контекст

Padding-маска стоит отдельного упоминания как источник ошибок: если её забыть, модель будет внимать заполнителям, и результат станет зависеть от того, как батч был собран. Ошибка не падает — просто качество на коротких последовательностях оказывается хуже, чем на длинных, и причина неочевидна.

Вариации по числу наборов ключей

Из урока про KV-кэш (140) станет понятно, зачем это нужно, но структура относится сюда:

схема	запросов	ключей и значений
MHA	$h$ наборов	$h$ наборов
MQA	$h$ наборов	один набор на все головы
GQA	$h$ наборов	$g$ групп, $1 < g < h$

Вычислений это почти не меняет, а вот размер кэша делит на $h$ (для MQA) — и именно кэш, а не арифметика, ограничивает генерацию. Поэтому современные модели почти все используют GQA: компромисс между выразительностью MHA и памятью MQA.