Внимание и трансформер

Attention позволяет каждой позиции построить собственную взвешенную выборку из контекста. Запрос говорит, что мы ищем; ключ — что может предложить позиция; value — какой сигнал будет перенесён. Карта внимания не хранится в весах: она вычисляется заново для каждого текста, и потому одна и та же голова на разных фразах смотрит в разные стороны.

Местоимение ищет опору

В предложении «Робот положил стакан на стол, потому что он был свободен» местоимение связано со столом, а не со стаканом. Между словами несколько позиций, и фиксированного локального окна может не хватить: поменяйте «свободен» на «полон» — и опора переедет на стакан, хотя расстояния останутся теми же. Рекуррентная сеть должна пронести нужный факт через скрытое состояние, шаг за шагом, теряя его в произведении матриц. Attention поступает иначе: он строит прямое ребро от одной позиции к другой и делает это заново для каждой фразы.

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

Три матрицы: запрос, ключ, значение

Пусть матрица представлений $X\in\mathbb R^{n\times d}$ содержит по строке на позицию. Три обучаемые проекции дают три разные роли:

$$Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V.$$

Строка $q_i$ — query позиции $i$, строка $k_j$ — key позиции $j$, строка $v_j$ — переносимое value. Оценка совместимости и веса:

$$s_{ij}=\frac{q_i^\top k_j}{\sqrt{d_k}},\qquad a_{ij}=\frac{\exp s_{ij}}{\sum_{\ell=1}^{n}\exp s_{i\ell}},$$

а выход позиции — взвешенная сумма значений:

$$z_i=\sum_{j=1}^{n}a_{ij}v_j.$$

В матричной записи всё это одна формула:

$$\operatorname{Attn}(Q,K,V)= \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V .$$

Отсюда сразу следует геометрическое ограничение одной головы:

$$z_i\in\operatorname{conv}\{v_1,\ldots,v_n\}, \qquad \|z_i\|\le\max_j\|v_j\| .$$

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

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Бюджет внимания нельзя увеличить

Строка карты — распределение вероятностей. Это значит, что внимание — игра с нулевой суммой: чтобы одна позиция получила больше, другие обязаны получить меньше. Формально при изменении одного score

$$\frac{\partial a_i}{\partial s_i}=a_i(1-a_i),\qquad \frac{\partial a_i}{\partial s_j}=-a_ia_j\ \ (i\ne j).$$

В матричной записи якобиан softmax равен

$$\frac{\partial a}{\partial s}=\operatorname{diag}(a)-aa^\top,$$

и он вырожден: вектор из единиц лежит в его ядре, потому что прибавление константы ко всем scores ничего не меняет. Отрицательный знак вне диагонали — это и есть конкуренция. Отсюда следует неприятный практический вывод: вес зависит не только от своего score, но и от того, кто ещё стоит в очереди. Возьмём опору со score $4$ и конкурентов со score $1$. При одном конкуренте вес опоры равен $0{,}953$; добавим ещё десять одинаковых нерелевантных токенов — и он падает до $0{,}646$, а при сорока — до $0{,}329$. Score опоры не изменился ни разу; изменился знаменатель.

Чтобы говорить о «размытости» карты количественно, удобна энтропия строки

$$H(a_i)=-\sum_{j=1}^{n}a_{ij}\ln a_{ij}, \qquad 0\le H(a_i)\le\ln n,$$

где нуль отвечает жёсткому выбору одной позиции, а $\ln n$ — полному отказу выбирать.

Почему делим на корень размерности

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

$$\mathbb E[q_rk_r]=0,\qquad \operatorname{Var}(q_rk_r)=\mathbb E[q_r^2k_r^2]=1,$$

а для скалярного произведения

$$\operatorname{Var}\!\left(\sum_{r=1}^{d_k}q_rk_r\right)=d_k, \qquad \operatorname{Var}\!\left(\frac{q^\top k}{\sqrt{d_k}}\right)=1 .$$

Модельный эксперимент (seed 760, восемь ключей на запрос) подтверждает это с точностью до третьего знака: измеренная дисперсия при $d_k=64$ равна $64{,}2$, а отношение дисперсии к $d_k$ по всем размерностям от 4 до 512 лежит между $0{,}977$ и $1{,}003$. Последствия для softmax видны сразу. Без деления при $d_k=512$ средний максимальный вес равен $0{,}955$, а энтропия карты падает до $0{,}107$ при максимуме $\ln 8=2{,}079$: распределение почти one-hot. С делением максимальный вес остаётся около $0{,}360$ при любой размерности, а энтропия — $1{,}729$.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Это не косметическая константа, а нормировка масштаба — родня приёмам инициализации сетей: при росте числа слагаемых сумма обязана оставаться управляемой. Заметьте и обратную сторону: если бы мы хотели «острого» внимания, достаточно было бы поделить на меньшее число. Так и делают, когда вводят температуру $\tau$ и пишут $\operatorname{softmax}(s/\tau)$.

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

Карта внимания на реальном тексте

Абстракцию полезно приземлить. Возьмём корпус SMS Spam Collection: $5569$ сообщений, из них $747$ спамовых, то есть $13{,}4\%$. По самому корпусу построим словарь из $1045$ слов, встречающихся не реже десяти раз, посчитаем матрицу совместной встречаемости в окне четырёх слов, превратим её в PPMI и возьмём $32$ первых сингулярных направления. Получатся честные embeddings, выученные из этих же SMS — техника прямо из урока об ортогональности и PCA. Ближайшие соседи выглядят осмысленно: к слову free ближе всего nokia ($0{,}912$) и mobile ($0{,}850$), к слову call — land ($0{,}847$) и claim ($0{,}799$).

Теперь возьмём реальное сообщение из корпуса, оставим девять различных слов free nokia motorola on orange call mobileupd or optout и построим одну голову без обучения: пусть $W_Q=W_K=W_V$ будут пропорциональны единичной матрице, то есть ключ, запрос и значение — это сам embedding. Карта получается содержательной. Запрос call отдаёт себе вес $0{,}616$, а слову free — всего $0{,}036$: в этом корпусе call живёт рядом с claim и land, а не рядом с free. Энтропия строки равна $1{,}413$ при равномерном максимуме $\ln 9=2{,}197$ — карта заметно сфокусирована, но не вырождена.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Маска меняет множество допустимого

В декодере языковой модели позиция $i$ не должна видеть будущие токены. К scores прибавляют матрицу

$$M_{ij}= \begin{cases} 0,&j\le i,\\ -\infty,&j>i, \end{cases} \qquad A=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}+M\right).$$

Эквивалентная запись — softmax по разрешённому множеству $\mathcal A_i$:

$$a_{ij}=\frac{\exp s_{ij}}{\sum_{\ell\in\mathcal A_i}\exp s_{i\ell}}\, \mathbb 1[j\in\mathcal A_i].$$

После softmax запрещённые веса равны нулю точно, а не приблизительно: $e^{-\infty}=0$. Padding mask убирает фиктивные позиции, дополняющие короткие строки до общей длины батча. Смысл масок разный: causal защищает от утечки будущего, padding — от искусственного содержимого, которого в тексте нет.

При реализации $-\infty$ заменяют большим отрицательным числом. В float16 диапазон невелик, и слишком большая по модулю константа даёт NaN уже на стадии вычитания максимума; библиотечные ядра учитывают это отдельно.

Позиция не входит в формулу сама

Посмотрите на определение внимания ещё раз: в нём нет ни одного индекса позиции, кроме как через содержимое строк. Если переставить строки $X$ перестановкой $P$, то

$$\operatorname{Attn}(PX)=P\operatorname{Attn}(X),$$

то есть self-attention эквивариантен перестановке. Численно это проверяется мгновенно: на нашей карте максимальное расхождение между переставленной картой и картой переставленного входа равно $1{,}1\cdot10^{-16}$ — машинный ноль. Без позиционной информации модель различает множество токенов, но не их порядок, и фразы «стакан на столе» и «стол на стакане» для неё одинаковы.

Классическое синусоидальное кодирование прибавляет к embeddings

$$\begin{aligned} \mathrm{PE}_{(p,2i)}&=\sin\!\left(p/10000^{2i/d}\right),\\ \mathrm{PE}_{(p,2i+1)}&=\cos\!\left(p/10000^{2i/d}\right). \end{aligned}$$

Скалярное произведение двух таких векторов зависит в основном от разности позиций: при $d=64$ оно равно $32{,}00$ для нулевого сдвига, $30{,}92$ для сдвига в одну позицию и $22{,}41$ для сдвига в восемь. Это и делает код полезным: близкие позиции похожи, далёкие — нет.

Есть и точное свойство: сдвиг позиции действует на код линейно,

$$\mathrm{PE}_{p+\Delta}=R_\Delta\,\mathrm{PE}_p, \qquad R_\Delta=\bigoplus_i \begin{pmatrix}\cos\omega_i\Delta&-\sin\omega_i\Delta\\ \sin\omega_i\Delta&\cos\omega_i\Delta\end{pmatrix}, \qquad \omega_i=10000^{-2i/d},$$

то есть каждая пара координат просто поворачивается на угол, зависящий только от расстояния.

Относительные схемы (в том числе RoPE, поворачивающий query и key на угол, пропорциональный позиции) выражают «на три позиции раньше» напрямую и потому лучше переносятся на длины, которых не было в обучении. Но переносимость надо измерять, а не постулировать: поведение за обученной длиной — отдельный экспериментальный вопрос, родственный вопросу об экстраполяции из урока о базисных признаках.

Несколько голов — несколько вопросов

Multi-head attention выполняет $H$ независимых наборов проекций:

$$\operatorname{head}_h= \operatorname{softmax}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_h}}\right)V_h, \qquad d_h=\frac{d}{H},$$

а затем склеивает результаты и смешивает их выходной матрицей:

$$\operatorname{MHA}(X)= \operatorname{Concat}(\operatorname{head}_1,\ldots,\operatorname{head}_H)\,W_O .$$

Число параметров при этом не растёт: все головы вместе несут столько же весов, сколько одна «широкая»,

$$H\cdot 3\,d\,d_h+d^2=3d^2+d^2=4d^2 .$$

Зачем это нужно, видно на паре построенных вручную голов. Голова A задана позиционно: её score равен $-3|j-(i-1)|$, то есть «смотри на предыдущий токен». Средний вес на предыдущий токен у неё $0{,}911$. Голова B задана содержательно: score равен скалярному произведению embeddings без диагонали, и её самая сильная связь — пара mobileupd–motorola с весом $0{,}354$. Две головы задают контексту два разных вопроса и обе полезны.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

В обученных моделях специализация действительно наблюдается, но она статистическая, а не гарантированная. Головы дублируют друг друга, часть из них можно удалить почти без потери качества, а «голова номер 3» в одном запуске и в другом — разные объекты: если переставить головы и согласованно переставить блоки $W_O$, функция $\operatorname{MHA}$ не изменится вовсе.

Внимание до нейросетей: рассогласование по Соколову

Слово «внимание» пришло в машинное обучение из физиологии, и русская школа занималась им задолго до 2017 года. Евгений Николаевич Соколов, изучая ориентировочный рефлекс, предложил в конце 1950-х идею нервной модели стимула: нервная система хранит след предъявлявшегося раздражителя и сравнивает с ним каждый новый вход. Совпадение приводит к угасанию реакции, рассогласование — к вспышке ориентировочной реакции, то есть к повороту внимания на изменившийся элемент.

Структура рассуждения та же, что в формуле внимания: есть хранимая модель (ключ), есть текущий вход (запрос), есть операция сравнения и есть следствие сравнения — усиление или ослабление передачи сигнала. Соколовская схема объясняла угасание реакции на повторяющийся звук и её мгновенное возвращение при малейшем изменении тона; в терминах нашей формулы угасание — это падение score, а вспышка — его рост. Позже Соколов развил идею векторного кодирования: стимул представляется вектором активностей нейронов-детекторов, а близость стимулов — углом между векторами. Это буквально та геометрия, в которой мы считаем $q^\top k$.

Александр Романович Лурия в книге «Внимание и память» (1975) формулировал ту же мысль с другой стороны: внимание есть фактор избирательности, без которого всякий психический процесс тонет в потоке сигналов. Механизм внимания в сети устроен так же прагматично: ёмкость представления ограничена, и её приходится распределять.

Лаборатория внимания

Выберите запрос и двигайте ползунок «прибавка к score ключа „стол“». Следите не за поднятым столбиком, а за соседними: softmax перераспределяет всю массу, и чужие веса падают, хотя их scores не менялись. Затем нажмите «подменить value слова „стол“»: веса останутся теми же до последнего знака, а выход переедет — это и есть разделение адреса и содержимого.

Переключите масштаб logits в режим «без деления» и увеличьте $d_k$ до 256: распределение схлопнется в один столбик, максимальный вес подскочит, энтропия упадёт почти до нуля. Верните деление на $\sqrt{d_k}$ — картина восстановится при любой размерности. Наконец включите causal mask и попробуйте связать позицию с будущим: нулевой вес там — следствие запрета до softmax, а не результат обучения. Кнопка «добавить 6 шумовых токенов» показывает разбавление: вес опоры падает от одного роста знаменателя.

Обучаемое взвешивание: одна голова на реальном спаме

Проверим внимание в деле, на тех же SMS. Модель максимально простая: у нас есть фиксированные embeddings $e_w$ из корпуса и обучаемый вектор-запрос $u$. Для сообщения из токенов $w_1,\ldots,w_m$ считаем

$$s_j=\frac{u^\top e_{w_j}}{\sqrt d},\qquad a=\operatorname{softmax}(s),\qquad c=\sum_{j=1}^{m}a_je_{w_j},$$

а затем обычный логистический классификатор:

$$\widehat p=\sigma(v^\top c+b),\qquad \sigma(t)=\frac1{1+e^{-t}} .$$

При $u=0$ формула вырождается в обычное усреднение:

$$a_j=\frac{\exp 0}{\sum_{\ell=1}^{m}\exp 0}=\frac1m, \qquad c=\frac1m\sum_{j=1}^{m}e_{w_j}.$$

Это внимание в чистом виде: вместо среднего по токенам, где $a_j\equiv 1/m$, сеть сама решает, кого слушать. Градиент по $u$ проходит через softmax:

$$\frac{\partial L}{\partial s_j}= a_j\left(g_j-\sum_{\ell}a_\ell g_\ell\right),\qquad g_j=(\widehat p-y)\,v^\top e_{w_j},$$

и хорошо видно, что общий сдвиг всех $g_j$ ничего не меняет — обучается только относительный порядок токенов.

На $1368$ отложенных сообщениях среднее по токенам даёт точность $96{,}3\%$, одна обученная голова — $96{,}6\%$; константный ответ «не спам» дал бы $87{,}0\%$. Прибавка скромная, $0{,}3$ процентных пункта, то есть $51$ ошибка против $47$: на коротких SMS усреднение и так почти не мешает. Но на сообщениях средней длины (от 10 до 25 токенов, их $586$) разрыв заметнее: $94{,}5\%$ против $95{,}2\%$ — ровно там, где разбавление начинает вредить.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Стоит задержаться на этом списке. Он честно рассказывает, чем спам отличается от переписки в этом корпусе: адресами, доменами, тарифами, словом txt. И он же показывает границу метода: голова научилась не «понимать спам», а находить лексические маркеры именно этой коллекции 2000-х годов — ровно та опасность, о которой шла речь в уроке про утечки и валидацию.

Квадратичная цена

Матрица $QK^\top$ имеет $n^2$ элементов, поэтому время и память одной головы растут как

$$T(n)=O(n^2d_k),\qquad M(n)=O(n^2).$$

Для блока из $H$ голов размерности $d_h=d/H$ это даёт

$$T_{\text{MHA}}(n)=O\!\left(Hn^2d_h+nd^2\right)=O\!\left(n^2d+nd^2\right),$$

и первое слагаемое обгоняет второе, как только $n$ превосходит $d$.

Это не теоретическая страшилка, а измеримый факт. Прямой замер на одной машине (numpy, float32, $d_k=64$, медиана из пяти повторов) даёт наклон в логарифмических координатах $b\approx1{,}9$: удвоение длины контекста с 512 до 1024 увеличило время примерно вчетверо. Абсолютные миллисекунды у вас будут свои, показатель степени — тот же. Память карты считается совсем точно: при $n=2048$ во float32 это $16{,}0$ МиБ на голову на слой, при $n=8192$ — уже $256{,}0$ МиБ.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

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

Стоимость авторегрессивной генерации складывается по шагам:

$$\sum_{t=1}^{n}O(td_k)=O(n^2d_k)$$

с кэшем против $\sum_{t=1}^{n}O(t^2d_k)=O(n^3d_k)$ без него.

Cross-attention и выравнивание перевода

В схеме encoder–decoder запросы приходят из декодера, а ключи и значения — из энкодера:

$$Z=\operatorname{softmax}\!\left( \frac{Q_{\text{dec}}K_{\text{enc}}^\top}{\sqrt{d_k}}\right)V_{\text{enc}} .$$

Формула та же, разные только источники. На корпусах перевода такую карту сравнивают с человеческим выравниванием слов — и сразу упираются в то, что перевод не бывает один-к-одному: артикль не имеет русского соответствия, а целая фраза перестраивается. Поэтому оценивать стоит не только BLEU, но и конкретные явления: имена, числа, отрицания, дальние согласования. Карта внимания здесь — инструмент поиска гипотезы об ошибке, а не приговор: итоговый выход зависит ещё и от residual-пути, от MLP и от всех остальных слоёв трансформера.

Карта — не объяснение

Соблазн велик: нарисовать красивую тепловую карту и объявить её объяснением решения. Против этого есть и теоретические, и экспериментальные возражения.

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

Проверка делается вмешательством. Возьмите сто реальных текстов, выберите голову и токен с максимальным весом. Замаскируйте его, потом токен со средним весом, потом случайный, и сравните изменение нужного logit. Если маскирование максимального веса не влияет сильнее прочих, карта на этом наборе ничего не объясняет. У маскирования, впрочем, есть свой изъян: получается строка, которой не бывает в данных. Мягче работает activation patching — подмена value одного токена значением из контрольного примера при сохранении остального контекста.

Мини-исследование: задача адресации

Синтетическая задача проверяет механизм лучше любой heatmap. Генерируйте строки вида A\:7 B\:2 C\:9 ?B и требуйте вернуть значение запрошенного ключа. Меняйте число пар от 2 до 64, порядок и расстояние до запроса. Здесь известен правильный паттерн внимания: запрос обязан найти соответствующий ключ и перенести именно его value.

Обучите одну голову и проведите тройку вмешательств. Переставьте пары — ответ обязан сохраниться (эквивариантность). Замените ключ при прежнем value — адрес должен смениться. Замените value при прежнем ключе — вес должен остаться, а выход измениться. Эта тройка проверяет функциональное разделение $Q$, $K$ и $V$ строже, чем любая картинка.

Затем добавьте два одинаковых ключа с разными значениями и правило «брать последний». Без позиционной информации задача неразрешима в принципе: вход, с точностью до перестановки, один и тот же. С относительным кодированием модель может предпочесть ближайшее прошлое. Измерьте точность по длинам, превышающим обучающие, — и вы получите честную кривую переноса вместо общих слов.

Адресация не равна пониманию

Соберём урок в одну картину. Внимание — динамическая адресация контекста: запрос сравнивается с ключами, softmax превращает совместимость в конкурирующие веса, и выход собирается из values, оставаясь в их выпуклой оболочке. Деление на $\sqrt{d_k}$ удерживает масштаб logits, а с ним и градиенты. Маска задаёт множество допустимых связей. Позиционное кодирование возвращает порядок, которого в формуле нет. Головы задают несколько вопросов параллельно. Всё это стоит $n^2$ и упирается в длину контекста.

И главное ограничение — методологическое. Карта показывает, откуда смешано содержимое, а не почему модель ответила именно так. Рисунок внимания хорош, когда отвечает на конкретный вопрос и проверяется вмешательством; во всех остальных случаях это цветной орнамент. Следующий шаг — увидеть, как этот механизм встраивается в блок с residual и нормализацией и превращается в трансформер.

Задачи