#Внимание

После эмбеддинга каждый токен — это вектор, но каждый вектор знает только о своём токене. Вектор слова он одинаков в любом предложении — неважно, имеется в виду кот, проспект или погода. Чтобы хорошо предсказывать, каждый токен должен подтянуть информацию из других токенов. Механизм, который это делает, называется вниманием (attention). Именно эта идея заставила трансформеры работать.

#Интуиция

Прочитайте предложение:

Кот не перешёл проспект, потому что он был слишком уставшим.

Кто был уставшим? Кот. Теперь замените одно слово:

Кот не перешёл проспект, потому что он был слишком широким.

Теперь «он» — это проспект. Смысл «он» определяют слова вокруг.

Значит, каждый токен должен смотреть на другие. Но не на все одинаково: чтобы понять уставшим, слово Кот очень важно, а что — почти нет. Простое среднее всех предыдущих векторов утопило бы полезное слово в шуме. И какие слова важны, зависит от того, кто спрашивает: уставшим нужен кто-то живой, широким — какое-то место.

Внимание решает эту задачу. Каждый токен даёт каждому предыдущему токену (и себе самому) вес — его долю внимания, и в сумме веса дают 100%. Новый вектор токена — смесь векторов, на которые он смотрел, в этих пропорциях.

Рисунок 5.1 · Куда смотрит одна голова внимания

Кто такой «он»?Предыдущее словоПоиск подлежащегоОтдых

…уставшим…широким

Голова кореференции: последнее слово смотрит назад, чтобы понять, к кому относится «он». Наведите курсор, нажмите или используйте стрелки на слове слева. Точки показывают, как информация течёт от слов, на которые оно смотрит; серые слова — его будущее, они скрыты маской.

Обратите внимание на одну честную деталь. В модели типа GPT слово он ещё не видит слова уставшим: оно идёт позже. Поэтому на позиции он голова колеблется между Кот и проспект. Неоднозначность снимается на последнем слове, которое может оглянуться на всё предложение. Информация в тексте течёт только вперёд.

#Запросы, ключи и значения

Откуда берутся веса? Модель должна для каждой пары токенов решить, насколько один важен для другого, и у неё нет ничего, кроме их векторов. Она делает это так, как вы ищете книги в библиотеке:

• Вы приходите с запросом (query): «книги про кошек».
• На корешке каждой книги есть ярлык — ключ (key): «домашние животные», «кулинария», «Париж».
• Вы сравниваете запрос с каждым ярлыком и уносите домой содержимое подошедших книг — их значения (value).

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

• запрос говорит, что токен ищет (уставшим: «кого-то, кто может устать»),
• ключ говорит, что токен предлагает тем, кто ищет (Кот: «я живое существо»; проспект: «я место»),
• значение — информация, которую токен отдаёт, когда его выбрали.

Ключи и значения разделены, потому что то, по чему токен легко найти, — не то же самое, что полезно передать.

#Маленький пример с числами

Пусть в векторах всего два числа: первое означает «это живое существо», второе — «это место», как в анкете из главы 3. Запрос слова уставшим — $q = (2, 0)$: ему нужно что-то живое. Он видит три ключа, включая собственный:

Оценка — это скалярное произведение: перемножаем числа на одинаковых местах и складываем, так что Кот получает $2 \times 1{,}5 + 0 \times 0 = 3$. Оно велико, когда запрос и ключ смотрят в одну сторону. Softmax (на входе оценки, на выходе вероятности) превращает оценки $3, 0, 0$ в веса примерно $91\%$, $4{,}5\%$ и $4{,}5\%$. Новый вектор слова уставшим:

$$\begin{aligned} &0{,}91\, v_\text{кот} \\ +\ &0{,}045\, v_\text{проспект} \\ +\ &0{,}045\, v_\text{уставшим} \end{aligned}$$

почти чистая информация о коте. Дайте последнему слову запрос $q = (0, 2)$, как у широким («что-то, что может быть широким»), и те же ключи отдадут $91\%$ проспекту. Те же слова, другой запрос — другая смесь.

В настоящей модели никто не решает, что означают числа. Каждый токен получает три своих вектора, умножая собственный вектор $x$ на три обучаемые матрицы:

$$q = x W_Q \qquad k = x W_K \qquad v = x W_V$$

а обучение подправляет $W_Q$, $W_K$ и $W_V$, пока получающиеся совпадения не начнут помогать предсказывать следующий токен. В векторах десятки или сотни чисел, и у их направлений редко бывают такие чистые имена, как «живое существо» (здесь это упрощение, чтобы всё можно было посчитать в уме).

#Все токены сразу

Каждый токен проводит такой поиск одновременно с остальными, по токенам своего же текста, поэтому это называется самовниманием (self-attention). Сложите запросы в матрицу $Q$ (по строке на токен), ключи — в $K$, значения — в $V$, и всё вычисление записывается одной формулой:

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V$$

Читайте её изнутри наружу:

• $K^\top$ — это транспонированная матрица $K$: её строки стали столбцами, поэтому $QK^\top$ — таблица сразу всех оценок «запрос · ключ»: в строке $i$, столбце $j$ стоит то, насколько запрос токена $i$ совпадает с ключом токена $j$.
• Деление на $\sqrt{d_k}$, где $d_k$ — длина каждого вектора-ключа, не даёт оценкам стать слишком большими (почему — ниже).
• $M$ — каузальная маска, то есть правило «токен видит только себя и предыдущие токены»: $0$, если токен $j$ стоит не позже токена $i$, и $-\infty$ для будущего. Так как $e^{-\infty} = 0$, будущие токены получают ровно нулевой вес.
• Softmax превращает каждую строку оценок в веса с суммой 1, а умножение на $V$ смешивает значения с этими весами.

Зачем маска? Модель пишет текст слева направо, и когда она выдаёт слово, следующих слов ещё не существует. При обучении весь текст обрабатывается сразу, и маска не даёт каждой позиции подглядеть в те самые слова, которые ей нужно предсказать.

Теперь то же самое по-настоящему, шаг за шагом: четыре токена по четыре числа, одна голова и, для простоты, без поворота позиций. Нажмите «пуск» или шагайте стрелками. Каждое число на экране считается вживую из предыдущих, так что на шаге с оценками выделенную клетку можно проверить вручную.

Рисунок 5.2 · Самовнимание, шаг за шагом

X · векторы токенов

d₀d₁d₂d₃

Кот0.2−0.40.10.9

·сел1.00.3−0.60.2

·на−0.30.80.5−0.1

·пол0.40.60.9−0.5

1/12

Четыре токена, каждый уже превращён в вектор из 4 чисел (в настоящих моделях — тысячи). Сложенные вместе, они образуют матрицу X.

Стоит заметить три вещи:

1. Всё — это умножение матриц. Все запросы сравниваются со всеми ключами одним произведением. Поэтому внимание так хорошо работает на GPU (видеочипах с тысячами маленьких ядер, которые умножают параллельно), а обучение обрабатывает все позиции последовательности параллельно.
2. Выход той же формы, что и вход. Четыре вектора на входе — четыре на выходе, поэтому слои внимания можно ставить друг на друга.
3. Веса зависят от входа. Обычное умножение на матрицу, например на $W_Q$, применяет к любому входу одни и те же числа. Здесь веса смешивания $A$ вычисляются из самих токенов. Именно поэтому «он» может значить разное в разных предложениях.

#Зачем делить на √d?

Сложите $d$ произведений случайных чисел порядка $\pm 1$, и сумма по величине будет обычно около $\sqrt{d}$: скалярное произведение двух векторов из 4 чисел — около 2, из 1024 чисел — около 32. Для больших векторов «сырые» оценки становятся большими, а softmax от больших чисел почти «one-hot»: один ключ забирает почти 100%, градиенты для всех остальных (сигналы, по которым обучение их подправляет) почти исчезают, и обучение застревает.

Рисунок 5.3 · Насыщение softmax

Размер вектора d_k64

Скоры разбросаны на 27.4 (скалярное произведение d случайных чисел растёт как √d = 8).

softmax(q·k)

ключ 1100.0%

ключ 20.0%

ключ 30.0%

ключ 40.0%

ключ 50.0%

ключ 60.0%

softmax(q·k ÷ √d)

ключ 149.5%

ключ 21.6%

ключ 36.6%

ключ 416.3%

ключ 510.2%

ключ 615.7%

Деление на $\sqrt{d_k}$ держит оценки в районе единицы, какой бы широкой ни была голова.

#Много голов сразу

Один шаблон внимания выражает только один тип связи. У токена один набор весов с суммой 100%, и если слову уставшим нужны и существительное, к которому оно относится, и слово прямо перед ним, одной голове придётся разделить внимание и размыть оба. Поэтому трансформер запускает несколько голов параллельно, у каждой свои $W_Q$, $W_K$, $W_V$. При ширине модели $d_\text{model} = 4096$ и 32 головах каждая голова работает в своём 128-мерном срезе.

Рисунок 5.4 · Четыре головы, одно предложение

Кто такой «он»?

Кот

не

перешёл

проспект,

потому

что

он

был

слишком

уставшим

Предыдущее слово

Кот

не

перешёл

проспект,

потому

что

он

был

слишком

уставшим

Поиск подлежащего

Кот

не

перешёл

проспект,

потому

что

он

был

слишком

уставшим

Отдых

Кот

не

перешёл

проспект,

потому

что

он

был

слишком

уставшим

Строки — запросы, столбцы — ключи (в том же порядке слов). Наведите на клетку.

Выходы всех голов склеиваются обратно в один вектор ширины $d_\text{model}$ и умножаются ещё на одну обучаемую матрицу $W_O$, которая смешивает найденное головами:

$$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\, W_O$$

Исследователи, заглядывающие внутрь обученных моделей, действительно находят головы с узнаваемой работой: головы, смотрящие на предыдущий токен; индукционные головы (induction heads), которые находят более раннюю копию текущего токена и предсказывают то, что шло за ней в прошлый раз (основа копирования шаблонов из промпта); головы, связывающие глаголы с подлежащими; и много голов, которые «отдыхают» на первом токене. Но большинство голов делают вещи, которым трудно дать имя.

Головам не обязательно иметь собственные ключи и значения. Во внимании с общими ключами (multi-query attention, MQA) все головы запросов делят одну голову ключей и одну голову значений; в групповом внимании (grouped-query attention, GQA), которое используют большинство современных открытых моделей, одну голову ключей делит небольшая группа голов запросов. Запросы остаются разнообразными, а память под ключи и значения сжимается в несколько раз. Почему эта память так важна, показывает глава 10.

#Один полный проход от начала до конца

Вот всё, что один слой внимания делает с текстом из $n$ токенов, с размерами Llama 2 7B ($d_\text{model} = 4096$, 32 головы по 128 чисел):

1. Вход: $n$ векторов по 4096 чисел, по одному на токен.
2. Проекции: каждая голова умножает их на свои $W_Q$, $W_K$, $W_V$, и каждый токен получает запрос, ключ и значение по 128 чисел. С RoPE запросы и ключи именно здесь поворачиваются в зависимости от позиции.
3. Оценки: каждая голова сравнивает каждый запрос с каждым ключом — таблица $n \times n$ скалярных произведений, делённых на $\sqrt{128} \approx 11{,}3$.
4. Маска и softmax: будущее заменяется на $-\infty$, а softmax превращает каждую строку в веса с суммой 1.
5. Смешивание: выход токена в этой голове — взвешенное среднее видимых ему значений, 128 чисел.
6. Объединение: выходы 32 голов склеиваются ($32 \times 128 = 4096$) и умножаются на $W_O$.
7. Выход: снова $n$ векторов по 4096 чисел, той же формы, что и вход. Каждый из них — поправка для своего токена, несущая то, что он собрал у других; как она прибавляется обратно, показывает следующая глава.

Шаги 2–5 — ровно то, что показывал рисунок 5.2, только с одной головой и четырьмя числами.

#Цена внимания

Каждый токен сравнивается с каждым предыдущим, поэтому объём работы растёт как квадрат длины последовательности $n$: вдвое больше текста — вчетверо больше оценок. Контекст в 128 000 токенов означает около $8 \times 10^9$ оценок на голову на слой. Что это значит для длинных контекстов, разбирается в главе 12.

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

Рисунок 5.5 · Какие оценки вычисляются

Полное каузальноеСкользящее окноСток + окноРазреженное с шагом

окно w4

70скоров посчитано из 136 при полном каузальном внимании528,6 млнпри 128k токенов (w = 4096) из 8,6 млрд при полном каузальном внимании

Окно плюс самый первый токен. Именно этот «сток внимания» позволяет модели с окном оставаться стабильной на бесконечном потоке текста.

Наведите на строку, чтобы увидеть, на что смотрит этот токен.

• Внимание со скользящим окном (sliding window) позволяет каждому токену видеть только последние $w$ токенов (в Mistral 7B $w = 4096$). Информация из более далёкого прошлого всё равно доходит — косвенно: каждый слой расширяет охват ещё на $w$.
• Стоки внимания (attention sinks). Помните «отдыхающую» голову, которая паркует вес на первом токене? Softmax обязан куда-то положить вероятность, и модели учатся использовать первый токен как «свалку». Если кэш с окном выбросит этот токен, качество рушится. Если его сохранить (StreamingLLM), всё работает; как этим пользуются кэши, показано в главе 11.
• Гибридные схемы. Многие современные модели чередуют слои: большинство использует локальное окно, а каждый несколько слоёв — полное глобальное внимание. Так сохраняется почти всё качество за долю стоимости.

Эти шаблоны меняют, что вычисляется. Другие приёмы сохраняют полное внимание, но считают его быстрее или хранят меньше: FlashAttention, KV-кэш и групповые запросы (глава 10).

#Итоги

• Внимание позволяет каждому токену собирать информацию из предыдущих токенов с весами, которые зависят от содержания.
• Каждый токен получает запрос (что он ищет), ключ (что предлагает) и значение (что отдаёт), умножая свой вектор на обучаемые матрицы.
• Запросы сопоставляются с ключами скалярными произведениями; softmax превращает совпадения в веса; веса усредняют значения.
• Оценки делятся на $\sqrt{d_k}$, чтобы softmax не насыщался, а каузальная маска скрывает будущее.
• Многоголовое внимание запускает много таких поисков параллельно и смешивает их через $W_O$; GQA и MQA позволяют головам делить ключи и значения.
• Полное внимание стоит $O(n^2)$, то есть работа растёт как квадрат длины; скользящие окна, стоки внимания и разреженные шаблоны считают меньше оценок.

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