Механизм внимания в трансформерах, который уже много лет “is all we need”, оказывается не таким уже незаменимым.

Это, конечно, не очень правдивый заголовок, но зато привлекает внимание. На самом деле речь пойдет о различных попытках заменить или хотя бы облегчить механизм внимания в трансформерах. Он во многом обеспечил успех моделей последних лет, но он же остается одним из узких мест. Дело в квадратичной сложности вычислений, из-за которой приходится ограничивать размеры входной последовательности. Ну и вообще, несмотря на все преимущества внимания, хотелось бы, чтобы оно, как основной двигатель моделей, было более экономично.

Попытки и подходы к модификации внимания начали появляться, наверное, почти сразу после появления его самого. Их можно (очень условно) разделить на группы:

Модифицировать внимание, чтобы облегчить вычисления
Вообще уйти от внимания к чему-то другому

К методам первой группы относятся все методы sparse attention, то есть прореживания тензоров внимания для уменьшения количества операций (пример). Еще один метод, точнее целый класс методов, — Hyena Hierarchy. Его суть в рекуррентном применении двух операций: умножения на диагональную матрицу и свертки с ядром, совпадающим по размеру с исходными данными:

Можно не прореживать внимание, а квантовать. Этот подход реализуют SageAttention и SpargeAttention . Идея первого — 8-битная квантизация самого внимания, а не только линейных слоев. Квантовать слои внимания сложно, потому что внутри него матричные умножения и софтмакс. То есть просто привести к 8 битам можно, но ничего хорошо из этого не получится. На картинке ниже результат text-to-image модели Unidiffuser с прямым квантованием.

Подобранный авторами алгоритм универсален для любых моделей, его можно внедрять без дообучения. Вывод ускоряется в 2,1-2,7 раза по сравнению с FlashAttention2 и xformers без потери в качестве.

Второй метод, SpargeAttention, использует эту квантизацию и дополняет его sparse attention. Сначала модель предсказывает карту внимания, то есть заранее отбрасывает лишние части — на них в прямом смысле не нужно обращать внимания. Отдельные части матриц Q и K избирательно сжимаются до одного токена на основе схожести токенов внутри блока. Затем в процессе софтмакса некоторые произведения просто пропускаются, основываясь на разнице между глобальным и локальными максимумами. И после всего этого применятся квантизация как в SageAttention.

Теперь перейдем ко второй группе направлений. Здесь поговорим о недавней работе аспиранта из университета Южной Калифорнии. Идея его метода SPECTRE такая — вместо блока внимания использовать преобразование Фурье и обрабатывать уже полученный фурье-образ. В целом, всё это полностью соответствует исконной идее фурье-преобразований: если в каком-то пространстве операции слишком сложные, нужно перейти в другое пространство, где всё будет выглядеть и считаться легче. Главное, чтобы потом можно было вернуться в исходное пространство без потерь.

Итак, к матрице всех построенных эмбеддингов применяем быстрое преоборазование Фурье. Теперь мы попали в пространство с ортонормальным базисом и уже в нем вводим что-то вроде внимания, а именно — обучаемый фильтр или гейт диагонального вида в этом же базисе. Альтернативно можно использовать не диагональную, а почти диагональную матрицу. В любом случае, главное, что удается перейти к простому перемножению. После этого совершаем обратное преобразование и возвращаемся в исходное пространство. Итоговая сложность — всего O(L log L),

Важный плюс этого подхода в том, что он легко совмещается с существующими архитектурами. Нет необходимости проводить дополнительную оптимизацию или перекраивать модель, достаточно просто заменить слои внимания на модули SPECTRE. Новые параметры составят меньше 6% модели, а если обучаемые фильтры одинаковы для разных голов, то и меньше 3%.

Экспериментальная часть пока не очень обширна (работа в большей степени дает теоретическое обоснование), но вот какие результаты показывают авторы:

SPECTRE сравнивали с FlashAttention2 на 32000 входных токенах по двум метрикам: время между получением промпта и выдачей первого токена (TTFT）и задержка перед получением каждого последующего токена (TPOT). Ускорение заметное, но это не за счет качества — с перплексией (PPL) на бенчмарке PG-19 у SPECTRE тоже все в порядке:

Важная оговорка, здесь все запуски были на входной длине всего в 1000 токенов. Так что пока это скорее предварительное тестирование. Главное увидеть, что мы не слишком много теряем в угоду скорости. Кстати, в статье не приведен код, так что для желающих отдельно даем ссылку.

Попробуем теперь проанализировать этот простой но интересный подход. LLM живет в многомерном пространстве знаний, где каждый токен можно рассматривать как дискретный сигнал. Тогда мы можем по аналогии перейти в частотное пространство этих сигналов. Что мы в нем увидим? Видимо, какую-то их общую динамику, корреляцию и взаимодействие в разных контекстах. Если исходные сигналы не белый шум, значит в их частотном представлении отразятся некие доминирующие компоненты, периодические паттерны и другие закономерности.

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

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

Пессимистичный взгляд на замену или упрощение внимания тоже имеет место. Раз из множества попыток заменить или упростить внимание ни одна не оказалась по-настоящему успешной, то возможно, лучше и не искать добра от добра. А вы какой взгляд разделяете?

Спасибо за внимание

Больше наших обзоров AI‑статей на канале Pro AI.

Внимание! Механизм внимания в трансформерах оказывается не таким уже незаменимым