Глава 2.1 — Внимание — это всё, что вам нужно

Содержание

1. Продолжаем с того места, где закончилась рекуррентность
2. Масштабированное скалярное внимание: сам механизм
3. Зачем нужен коэффициент масштабирования: вывод
4. Многоголовое внимание: почему одной функции внимания недостаточно
5. Стеки энкодера и декодера
6. Три применения одного механизма: self-attention, маскированное self-attention и cross-attention
7. Настоящий выигрыш: параллелизм
8. Взгляд с точки зрения собеседования
9. Вопросы для самопроверки
10. Источники

1. Продолжаем с того места, где закончилась рекуррентность

Глава 1.5 завершилась острым вопросом. Внимание Бахданау и внимание Луонга уже показали, что декодер может оглядываться на каждое состояние энкодера и учиться, куда именно смотреть, вместо того чтобы сжимать всё исходное предложение в один вектор фиксированного размера. Это был настоящий прорыв, но в любой seq2seq-системе, построенной на этой идее, внимание оставалось вспомогательным механизмом, пристроенным к рекуррентной сети (RNN). Основную работу по представлению последовательности по-прежнему выполнял рекуррентный "хребет", а внимание было лишь более умным способом читать из него информацию. Вопрос, которым заканчивалась глава, был поставлен прямо: если внимания достаточно, чтобы декодер мог извлечь именно ту информацию, которая ему нужна, из любой точки исходного текста, зачем вообще сохранять рекуррентность? Что, если внимание — это не вспомогательный механизм, а весь механизм целиком: не то, что читает представления, а то, что их строит?

Именно на этот вопрос в 2017 году ответили Vaswani et al., представив Transformer. Их аргумент состоял не столько в том, что внимание даёт представления лучше, чем рекуррентность, — хотя часто это действительно так, — сколько в архитектурных и вычислительных соображениях: рекуррентность несёт в себе неустранимую последовательную зависимость — нельзя вычислить скрытое состояние в позиции t, не вычислив его сначала в позиции t-1, — и эта зависимость становится узким местом при обучении, причём тем более серьёзным, чем больше становятся последовательности и наборы данных. Если бы рекуррентность можно было заменить механизмом, позволяющим каждой позиции напрямую обращаться к любой другой позиции, без пошаговой цепочки зависимостей, то представления для всей последовательности можно было бы вычислять параллельно. Это единственное архитектурное решение, пожалуй, самое значимое в истории глубокого обучения применительно к языку, потому что именно оно сделало вычислительно возможным обучение моделей на тех масштабах данных и параметров, которые определяют эпоху больших языковых моделей (LLM). Всё остальное содержание этой главы служит именно этой идее.

2. Масштабированное скалярное внимание: сам механизм

Уберём RNN полностью и зададимся вопросом: как позволить токену посмотреть на все остальные токены последовательности и самостоятельно решить, какие из них важны? Ответ Transformer — это масштабированное скалярное внимание (scaled dot-product attention), и его стоит выстроить с нуля, а не просто привести формулу, потому что каждая деталь в нём присутствует не случайно.

Эмбеддинг каждого токена проецируется тремя отдельными обучаемыми линейными отображениями в три вектора: вектор запроса (query) $q$, вектор ключа (key) $k$ и вектор значения (value) $v$. Запрос представляет собой "что этот токен ищет". Ключ представляет собой "что этот токен может предложить, как метку для поиска". Значение представляет собой "что этот токен фактически передаёт как содержание, если на него обратили внимание". Это разделение на query/key/value — прямое заимствование из метафоры поисковых систем: представьте это как мягкую таблицу поиска, где запрос сравнивается с набором ключей, чтобы получить оценку релевантности для каждой записи, и эти оценки определяют, какая доля значения каждой записи будет извлечена.

Для заданного запроса $q_i$ в позиции $i$ его сходство с каждым ключом $k_j$ в последовательности вычисляется как скалярное произведение, $q_i \cdot k_j$, для всех $j$. Затем эти сырые оценки делятся на $\sqrt{d_k}$, где $d_k$ — размерность векторов ключей, и пропускаются через softmax, чтобы превратить их в распределение вероятностей по позициям. Каждый выход тогда представляет собой взвешенную сумму всех векторов значений с этими весами: $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$ где $Q$, $K$ и $V$ — матрицы, составленные из векторов запросов, ключей и значений для каждой позиции последовательности. Выход для позиции $i$ — это взвешенное среднее всех векторов значений в последовательности, где веса полностью определяются тем, насколько хорошо запрос позиции $i$ совпал с ключом каждой позиции. Здесь ничто не зависит от выхода позиции $i-1$, как это было бы в RNN — веса внимания для каждой позиции можно вычислять независимо и параллельно, имея полный набор матриц $Q$, $K$, $V$.

3. Зачем нужен коэффициент масштабирования: вывод

Слагаемое $\sqrt{d_k}$ выглядит как незначительная деталь реализации, но на самом деле оно решает конкретную, количественно измеримую числовую проблему, и именно такие детали отличают тех, кто вывел Transformer самостоятельно, от тех, кто просто им пользуется. Предположим, как это примерно верно при инициализации, что компоненты $q$ и $k$ — независимые случайные величины с нулевым средним и единичной дисперсией. Скалярное произведение $q \cdot k = \sum_{l=1}^{d_k} q_l k_l$ тогда представляет собой сумму $d_k$ независимых произведений, каждое из которых имеет нулевое среднее и единичную дисперсию (поскольку $\text{Var}(q_l k_l) = E[q_l^2]E[k_l^2] = 1$ при этих предположениях). Суммирование $d_k$ независимых одинаково распределённых слагаемых означает, что дисперсия суммы в $d_k$ раз больше дисперсии одного слагаемого, поэтому $\text{Var}(q \cdot k) = d_k$.

Это важно из-за того, что происходит дальше: скалярные произведения подаются в softmax. По мере роста $d_k$ сырые скалярные произведения растут по величине пропорционально $\sqrt{d_k}$ (стандартному отклонению), а это значит, что при больших $d_k$ некоторые из этих оценок могут быть большими по абсолютной величине. Softmax, применённый к входам с большой дисперсией, становится крайне "заострённым" — почти вся вероятностная масса сосредотачивается на том входе, который наибольший, а градиент во всех остальных местах устремляется к нулю, потому что якобиан softmax уменьшается по мере того, как его выход приближается к one-hot вектору. На практике это означает, что при больших $d_k$ механизм внимания при инициализации будет стремиться "схлопнуться" в почти детерминированный жёсткий выбор, а градиенты, идущие обратно через softmax к проекциям запросов и ключей, будут затухать. Деление скалярных произведений на $\sqrt{d_k}$ в точности компенсирует рост дисперсии — $\text{Var}(q\cdot k / \sqrt{d_k}) = 1$ независимо от $d_k$, — что сохраняет распределение входов softmax в устойчивом виде вне зависимости от того, насколько велика размерность ключей, и поддерживает полезный поток градиента через внимание во время обучения. Это небольшой элемент арифметики с непропорционально большим эффектом: без него увеличение размерности модели активно работало бы против стабильности обучения — прямо противоположно тому, что нужно архитектуре, чьё главное предназначение — масштабироваться.

4. Многоголовое внимание: почему одной функции внимания недостаточно

Единственная функция внимания, даже хорошо масштабированная, заставляет каждый токен объединять все свои "причины для обращения внимания" в одно общее взвешивание по последовательности. Но в предложении обычно одновременно стоит отслеживать множество разных типов отношений — одно может касаться синтаксической зависимости (к какому глаголу относится это существительное), другое — кореференции (на какое более раннее существительное ссылается это местоимение), третье — простой позиционной близости. Одно распределение softmax по последовательности может выразить лишь одно взвешивание за раз, а значит, одноголовый механизм внимания вынужден усреднять все эти сигналы вместе, потенциально стирая различия, которые было бы полезно сохранить раздельно.

Многоголовое внимание (multi-head attention) решает эту проблему, запуская несколько независимых операций внимания параллельно, каждую в своём подпространстве меньшей размерности, и позволяя модели самостоятельно выучить, на чём специализируется каждая из них. Конкретно, модель обучает отдельные матрицы проекций для каждой из $h$ голов, проецируя один и тот же вход в $h$ различных наборов запросов, ключей и значений, каждый размерности $d_k = d_{\text{model}}/h$. Каждая голова вычисляет своё масштабированное скалярное внимание независимо, порождая $h$ выходных векторов на позицию, которые затем конкатенируются обратно и пропускаются через одну финальную обучаемую линейную проекцию, возвращающую размерность к $d_{\text{model}}$. Поскольку каждая голова работает в своём подпространстве со своими обученными проекциями, разные головы свободны специализироваться — эмпирически, после обучения, отдельные головы Transformer часто демонстрируют явно интерпретируемое поведение: например, одна голова стабильно обращает внимание на предыдущий токен, другая отслеживает согласование подлежащего и сказуемого, третья обращает внимание на первый токен последовательности как своего рода "холостой ход" или значение по умолчанию. Суть не в том, что какая-то отдельная голова сама по себе мощнее той единой большой функции внимания, которую она заменяет; суть в том, что предоставление модели нескольких независимых подпространств внимания меньшей размерности даёт ей больше представительной гибкости для одновременного отслеживания разных типов отношений, при примерно том же суммарном вычислительном затрате, что и одна операция внимания полной размерности.

5. Стеки энкодера и декодера

Исходный Transformer, представленный для машинного перевода, сохраняет структуру энкодер-декодер, уже установившуюся в seq2seq-моделях, но перестраивает обе половины полностью на основе внимания и простых полносвязных слоёв, без какой-либо рекуррентности. Энкодер представляет собой стек идентичных слоёв, каждый из которых содержит два подслоя: подслой многоголового self-attention, где каждая позиция обращает внимание на каждую позицию той же (исходной) последовательности, и позиционную полносвязную сеть (feed-forward network), которая применяется одинаково и независимо к представлению каждой позиции — по сути, небольшой двухслойный перцептрон с нелинейностью между слоями, запускаемый токен за токеном. Каждый подслой обёрнут остаточным соединением (residual connection), за которым следует нормализация слоя (layer normalization) — решение, полные последствия которого рассматриваются в главе 2.4.

Декодер повторяет эту структуру, но с дополнительным подслоем и важным ограничением. Каждый слой декодера содержит маскированное self-attention по уже сгенерированной части целевой последовательности, затем cross-attention, где позиции декодера служат запросами по отношению к выходу энкодера, выступающему в роли ключей и значений, и, наконец, тот же вид полносвязного подслоя. Стек из шести таких слоёв энкодера и шести слоёв декодера (в базовой конфигурации оригинальной статьи) даёт модель, в которой информация из исходного текста проникает в каждую позицию декодера через cross-attention на каждом слое, а целевая последовательность строит собственное внутреннее представление того, что уже было сгенерировано, посредством маскированного self-attention. В этом стеке нет последовательной зависимости между входными позициями энкодера; последовательная зависимость возникает только в декодере, и только потому, что генерация по своей природе последовательна — нельзя обратить внимание на токен, который ещё не сгенерирован.

6. Три применения одного механизма: self-attention, маскированное self-attention и cross-attention

Стоит чётко зафиксировать тот факт, что Transformer на самом деле не содержит трёх разных механизмов — он содержит один механизм, масштабированное скалярное внимание, используемый в трёх разных конфигурациях, и смешивание этих понятий — распространённый источник путаницы. Self-attention, применяемое в энкодере, позволяет каждой позиции обращать внимание на каждую позицию той же последовательности, включая позиции после неё — здесь нет никаких ограничений, потому что вся входная последовательность доступна сразу, и понятие причинности сохранять не нужно. Именно это делает представления энкодера двунаправленными: представление слова строится из контекста с обеих сторон.

Маскированное (каузальное) self-attention, используемое в декодере, — это та же операция с одной модификацией: перед softmax позиции маскируются так, что позиция $i$ может обращать внимание только на позиции $j \le i$. Обычно это реализуется путём добавления $-\infty$ к оценкам внимания для всех $j > i$ перед softmax, что приводит их вероятность после softmax к нулю. Именно эта маскировка сохраняет авторегрессивное свойство, необходимое для генерации — во время инференса токены генерируются по одному, и будущих токенов, на которые можно было бы обратить внимание, действительно не существует, поэтому обучение должно уважать то же самое ограничение, иначе модель научилась бы полагаться на информацию, которой у неё не будет во время инференса. Cross-attention, используемое во втором подслое декодера, — это снова тот же самый механизм, но с запросами, поступающими из собственной последовательности декодера, и ключами и значениями, поступающими из выхода энкодера — это прямой потомок внимания Бахданау и Луонга из главы 1.5, теперь выраженный как частный случай той же общей операции, а не как отдельный специально сконструированный механизм. Осознание того, что все три — это одна операция с разными входами, а не три отдельные идеи, — именно тот объединяющий взгляд, который должно порождать глубокое понимание этой архитектуры.

7. Настоящий выигрыш: параллелизм

Объясняя, почему победил Transformer, соблазнительно сосредоточиться исключительно на качестве представлений, но исторически решающим преимуществом было вычислительное, и об этом стоит сказать прямо. В RNN вычисление скрытого состояния в каждой позиции требует, чтобы скрытое состояние в предыдущей позиции уже было вычислено — граф вычислений имеет последовательную зависимость длиной, равной длине последовательности, и никакой дополнительный параллелизм оборудования не способен сократить этот критический путь, потому что каждый шаг действительно нуждается в результате предыдущего шага. Поэтому обучение на последовательности длины $n$ требует $O(n)$ последовательных операций, которые невозможно распараллелить по измерению последовательности, даже если в вашем распоряжении простаивает множество GPU.

У self-attention нет такой зависимости между позициями последовательности в рамках вычислений одного слоя. Для данного слоя запрос, ключ и значение для каждой позиции можно вычислить независимо, все попарные скалярные произведения можно вычислить как одно матричное умножение ($QK^\top$), и результирующие взвешенные суммы точно так же можно вычислить как одно матричное умножение. Это означает, что всё вычисление self-attention для последовательности длины $n$ можно выполнить как небольшое число крупных матричных умножений — а это ровно тот тип нагрузки, для параллельного "поглощения" которой построено современное оборудование GPU и TPU. Длина пути между любыми двумя позициями в сети теперь также составляет $O(1)$ вместо $O(n)$ — информация из позиции 1 может напрямую повлиять на представление позиции $n$ через одну операцию внимания, вместо того чтобы распространяться через $n$ последовательных рекуррентных шагов, что также, как правило, упрощает оптимизацию, поскольку градиентам не нужно "выживать" через длинную цепочку последовательных преобразований. Это конкретная, измеримая причина, по которой Transformer удалось обучать на наборах данных и в масштабах, непрактичных для RNN: не потому, что внимание — более остроумная идея в абстрактном смысле, а потому, что оно превращает принципиально последовательное вычисление в принципиально параллельное, а параллельные вычисления — это именно то, в чём хороши современные ускорители.

Впрочем, ничего из этого не даётся даром — self-attention в Transformer вообще не имеет встроенного понятия порядка. Поскольку внимание обрабатывает свой вход как множество векторов, а взвешенная сумма в формуле внимания инвариантна к порядку, в котором эти векторы перечислены, Transformer без каких-либо дополнительных модификаций буквально не способен отличить "собака укусила человека" от "человек укусил собаку" по одним лишь представлениям слов, поскольку перестановка входных позиций оставляет любое вычисление внимания неизменным. Решение этой проблемы — предоставление принципиально нечувствительному к порядку, полностью параллельному механизму способа узнать, где именно в последовательности находится каждый токен, — это ровно то, с чего начинается глава 2.2.

8. Взгляд с точки зрения собеседования

"Выведите, почему оценки внимания делятся на квадратный корень из размерности ключа. Что пошло бы не так без этого?" Сильный ответ напрямую проводит аргумент через дисперсию: предполагая единичную дисперсию, нулевое среднее и независимость компонент запроса и ключа, скалярное произведение двух $d_k$-мерных векторов имеет дисперсию $d_k$, так как это сумма $d_k$ независимых слагаемых, каждое с дисперсией 1. Без коррекции эта дисперсия растёт вместе с размерностью модели, толкая входы softmax к крайним значениям и вызывая насыщение softmax — приводя к почти one-hot выходам и затухающим градиентам через веса внимания. Деление на $\sqrt{d_k}$ нормализует дисперсию обратно к 1 независимо от размерности, удерживая softmax в хорошо обусловленном режиме. Самые сильные ответы отмечают, что это именно аргумент нормализации дисперсии, а не произвольная константа, и связывают это с тем, что значимость этого решения растёт, а не падает, по мере масштабирования моделей к большим значениям $d_k$.

"Почему используются несколько голов внимания вместо одной большей операции внимания?" Сильный ответ избегает расплывчатого невнятного "чтобы смотреть на разные вещи" и вместо этого объясняет, что многоголовое внимание разбивает представление на $h$ подпространств меньшей размерности, каждое со своими обученными проекциями query/key/value, так что разные головы могут специализироваться на разных типах отношений (синтаксических, позиционных, кореференциальных), не будучи вынужденными усреднять их в одно взвешивание. Также стоит упомянуть аспект вычислительной стоимости: разбиение на $h$ голов размерности $d_{\text{model}}/h$ сохраняет суммарные вычисления примерно сопоставимыми с одной головой полной размерности, так что это ближе к "перераспределению" ёмкости, чем к её "добавлению".

"Является ли self-attention инвариантным к перестановкам? Почему это важно?" Да — ответ должен не просто утверждать это, а показывать: поскольку выход внимания в позиции $i$ — это взвешенная сумма векторов значений с весами из softmax по скалярным произведениям, перестановка порядка входных токенов переставляет то, какое значение соответствует какому токену, но даёт совершенно тот же набор скалярных произведений и весов, так что у механизма самого по себе нет способа различить порядок последовательности. Именно поэтому позиционную информацию нужно вносить явно — это не мелкое упущение, а прямое структурное следствие того, что последовательность для целей внимания трактуется как неупорядоченное множество векторов.

"Почему устранение рекуррентности имело такое большое значение на практике, помимо качества представлений?" Сильный ответ сосредоточен на параллелизме: RNN несут в себе неустранимую $O(n)$-последовательную зависимость в графе вычислений, которую никакое параллельное оборудование не способно сократить, тогда как ядро вычисления self-attention — это несколько крупных матричных умножений без межпозиционной последовательной зависимости внутри слоя, что эффективно отображается на оборудование GPU/TPU. Кандидату также стоит упомянуть длину пути $O(1)$ между любыми двумя позициями против $O(n)$ у RNN, что также, как правило, облегчает оптимизацию.

"Разберите разницу между self-attention, маскированным self-attention и cross-attention." Ответ должен чётко утверждать, что это одна и та же операция масштабированного скалярного внимания с разными источниками для запросов, ключей и значений и разной маскировкой. Self-attention энкодера: Q, K, V все берутся из одной и той же (исходной) последовательности, без маскировки, полностью двунаправленно. Маскированное self-attention декодера: Q, K, V берутся из одной и той же (целевой) последовательности, при этом будущие позиции маскируются (добавлением $-\infty$ перед softmax), чтобы сохранить авторегрессивную генерацию. Cross-attention: Q поступает из декодера, K и V — из выхода энкодера, позволяя каждой позиции декодера читать из полного представления исходного текста.

9. Вопросы для самопроверки

1. Отталкиваясь от разложения по правилу цепочки, установленного в главе 1.1, и от постановки энкодер-декодер из главы 1.4, объясните своими словами, почему механизм, позволяющий токену напрямую обращать внимание на любой другой токен, — это естественный способ вычислять $P(x_i \mid x_1, \ldots, x_{i-1})$ без рекуррентности.
2. Выведите дисперсию $q \cdot k$ при стандартных предположениях инициализации и точно объясните, как деление на $\sqrt{d_k}$ устраняет проблему насыщения softmax.
3. Если бы вы увеличили $d_k$, не меняя коэффициент масштабирования, что бы вы ожидали увидеть в распределениях весов внимания на раннем этапе обучения и почему?
4. Объясните разделение на query/key/value с помощью аналогии с таблицей поиска и объясните, какую роль играет каждая из трёх обученных матриц проекций.
5. Какое именно вычислительное свойство self-attention отвечает за возможность параллельного обучения по позициям последовательности, и какое соответствующее свойство RNN оно заменяет?
6. Почему self-attention инвариантно к перестановкам и что из этого следует для любой полной архитектуры Transformer (то есть чего в ней не хватает)?
7. Опишите, чем маскированное self-attention механически отличается от немаскированного self-attention, и объясните, почему именно это отличие необходимо для авторегрессивной генерации.

10. Источники

• Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762. https://arxiv.org/abs/1706.03762