Внимание

Формула

$$\text{Attn}(Q,K,V) = \htmlData{k=weights}{\text{softmax}\!\left(\frac{\htmlData{k=scores}{QK^\top}}{\htmlData{k=scale}{\sqrt{d}}}\right)} V$$

Три шага, каждый из которых — контракция из урока 040:

$\texttt{'bhqd,bhkd->bhqk'} \quad\to\quad \text{softmax по } k \quad\to\quad \texttt{'bhqk,bhkd->bhqd'}$

Первый сворачивает размер головы, давая матрицу «каждый запрос против каждого ключа». Последний сворачивает ключи, возвращая размер головы. Обе операции стоят одинаково — внимание состоит из двух равно дорогих умножений, а не из одного.

Откуда √d

Это единственное место формулы, которое обычно принимают на веру. Аргумент чисто дисперсионный и проверяется в три строки.

Пусть компоненты $q$ и $k$ независимы, с нулевым средним и единичной дисперсией. Тогда

$q \cdot k = \sum_{i=1}^{d} q_i k_i, \qquad \operatorname{Var}(q \cdot k) = \sum_{i=1}^d \operatorname{Var}(q_i k_i) = d$

потому что дисперсия произведения независимых центрированных величин с единичной дисперсией равна единице, а дисперсии независимых слагаемых складываются (блок 3, урок 060).

Проверено численно:

$d$	измеренная дисперсия	предсказанная	ст. откл.
$4$	$4.03$	$4$	$2.01$
$64$	$64.21$	$64$	$8.01$
$1024$	$1017.8$	$1024$	$31.9$

Значит, без деления оценки имеют стандартное отклонение $\sqrt{d}$: при $d = 64$ логиты разбросаны на $\pm 8$, при $d = 1024$ — на $\pm 32$. Деление на $\sqrt{d}$ возвращает дисперсию к единице независимо от $d$. В этом весь смысл: масштаб не должен зависеть от размера головы.

Почему это важно, а не косметично

Если бы softmax был линейным, масштаб не имел бы значения. Но он насыщается, и последствия из урока 060 применяются напрямую.

Выключите масштабирование в виджете и посмотрите на два числа: энтропию строки и масштаб градиента.

делить на √d

ключи

запросы

	кот	сел	на	коврик	и	уснул
кот
сел
на
коврик
и
уснул

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

64

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

1

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

4.38

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

1.986 бит

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

0.8266

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

2.48e-1

Что стоит проделать:

1. выключите масштаб при $d = 64$ — энтропия строки падает, максимальный вес растёт;
2. поднимите $d$ до $256$ без масштабирования — насыщение усиливается, и появляется предупреждение про градиент;
3. включите масштаб обратно — распределение возвращается к осмысленному при любом $d$.

Числа для одного прогона при $d = 64$ и восьми ключах:

	разброс оценок	энтропия строки	максимальный вес
без масштаба	$25.1$	$0.994$ бита	$0.547$
с масштабом	$3.14$	$2.295$ бита	$0.355$

Равномерное распределение по восьми ключам дало бы $3$ бита. То есть без масштабирования внимание уже почти определилось, не увидев ни одного обучающего примера — просто из-за геометрии случайной инициализации.

Цена: где на самом деле уходят вычисления

Здесь стоит поправить распространённое утверждение. «Внимание квадратично, поэтому оно доминирует» — верно не всегда, и порог считается точно.

Для одного слоя с $h$ головами размера $d$ и моделью ширины $D = hd$:

часть	стоимость
проекции $Q, K, V, O$	$4 n D^2$
две контракции внимания	$2 n^2 D$

Отношение равно $\frac{n}{2D}$, то есть внимание перевешивает проекции только при $n > 2D$. Проверим при $D = 768$:

$n$	внимание / проекции
$128$	$0.08$
$1024$	$0.67$
$8192$	$5.3$

При коротком контексте внимание — малая часть работы, и оптимизировать его бессмысленно; основное время уходит в матрицы проекций. Квадратичный член начинает доминировать где-то за тысячей токенов, и вот тогда появляются FlashAttention, разреженные схемы и всё прочее.

Посмотрите на обе контракции в нотации индексов — квадратичность видна как две свободные буквы вместо одной.

оценкивыходпроекция

bhqd, bhkd → bhqk

индекс	размер	роль
b	1	остаётся
h	12	остаётся
q	64	остаётся
d	64	суммируется
k	64	остаётся

форма результата

[1, 12, 64, 64]

элементов

49,152

умножений-сложений

3,145,728

операций на элемент

64

Что ещё стоит знать

Внимание — это не поиск по ключам. Softmax даёт распределение по всем ключам, а не выбор одного. Название «query/key/value» пришло из аналогии со словарём, но операция мягкая: всегда используются все значения, просто с разными весами.

Оно инвариантно к перестановке. Формула не содержит информации о порядке токенов — переставьте их, и веса переставятся вместе с ними. Отсюда необходимость позиционных кодировок, о которых урок 110. Это не деталь реализации, а свойство самой операции.

Оно не имеет параметров. В формуле выше нет обучаемых весов вовсе — они все в проекциях $Q$, $K$, $V$, $O$. Само внимание — фиксированная функция трёх аргументов, и это одна из причин, по которой оно так хорошо переносится между задачами.