Как свергнуть короля: формула для линейной альтернативы Трансформеру без побочных эффектов

Привет! На связи Андрей Аргаткин, руководитель научной группы исследований эффективных архитектур нейронных сетей ИМШ ВШЭ. Я хочу рассказать о нашем текущем исследовании в рамках совместного образовательного проекта с VK. В ходе исследования мы надеемся выделить волшебную формулу из недавней архитектуры DANet (1, 2) и экстраполировать её на широкий спектр других моделей, что позволит им стать такими же крутыми по качеству, но гораздо более быстрыми и эффективными, чем бессменный король мира нейронных сетей — Трансформер. Эта формула должна избавить от побочных эффектов предыдущих архитектур, пытавшихся стать ему заменой. Но сначала поговорим, зачем всё это вообще нужно.

Проблема Трансформера, с которой все (почти) смирились

Доминирующей архитектурой для обработки последовательностей до сих пор, спустя 8 лет после своего создания, остаётся Трансформер. Его используют на практике не только в больших языковых моделях, но и для других модальностей, которые можно представить в виде последовательностей — обработки изображений, звука, видео, временных рядов, данных генома и т. д. Однако у стандартных трансформеров есть ахиллесова пята: их вычислительная сложность растёт квадратично ($O(N^2)$) с увеличением длины последовательности N. Это делает затратной по стоимости и долгой по времени обработку длинных текстов, документов, видео и других данных, а также приводит, например, к лимитам на количество запросов или размер контекстного окна в интерфейсах ChatGPT и других чат-ботов по сравнению с тем, что заявлено разработчиками моделей и доступно платящим за API пользователям. 

Способность эффективно и быстро обрабатывать очень длинные контексты была бы очень кстати для рассуждающих языковых моделей, таких как OpenAI o3 или DeepSeek R1, которые показывают великолепные результаты на самых разных задачах благодаря длинным цепочкам рассуждений (chains of thought). Качество ответа у таких моделей напрямую зависит от «бюджета» в токенах, отведённого ей на размышления: чем дольше она сможет думать, тем точнее результаты. Этот простой принцип лежит в основе очень эффективной парадигмы test-time scaling — масштабирования во время инференса, — которая менее чем за год позволила OpenAI пройти путь от o1 до o4 и сравняться с людьми на многих недосягаемых ранее бенчмарках. 

Однако такие достижения требуют огромных затрат: чтобы улучшить результат всего на 12 % (с 76 % до 88 %) на соревновании ARC-AGI модели o3 потребовалось в 172 раза (172, Карл!) больше вычислений, с примерно $200 до $34 000 долларов за каждую из сотен задач. При этом из таблицы в блог-посте следует, что на каждую из 1024 попыток решения одной задачи тратилось в среднем по 55,5 тысяч токенов рассуждений, или $33,5 в денежном эквиваленте.

Этот пример наглядно показывает, что нам не помешали бы другие, вычислительно более эффективные альтернативы Трансформеру, имеющие лучшую, чем  $O(N^2)$, временную асимптотику. 

Две волны попыток свергнуть короля

Линейные трансформеры

За прошедшие годы было предложено много новых архитектур и модификаций Трансформера, призванных улучшить вычислительную сложность стандартного механизма softmax-внимания (softmax attention). Я условно разделяю их на две волны. Модели первой волны (2019-2022 годы) видоизменяли или аппроксимировали механизм внимания (ссылки на Transformers are RNNs, Reformer, Performer, Transformer quality in linear time). Если в стандартном Трансформере он считается (в упрощённом виде) как: 

то большинство новых архитектур можно свести к обобщённой формуле:

где , и — некоторые функции (в случае Трансформера и многих других архитектур это линейные проекции) от последовательности входных эмбеддингов , а — некоторое преобразование (ядерная функция), которое делает элементы матриц , после применения неотрицательными. Диагональная матрица предназначена для того, чтобы построчно нормировать матрицу-произведение так, чтобы сумма элементов каждой строки была равна 1. Это свойство и неотрицательность элементов повторяют свойства функции softmax и делают алгоритм вычислительно стабильным, а благодаря тщательно подобранной форме функции результат линейного внимания можно сделать приближённо похожим на softmax-внимание. 

Так как матрицы и в формуле линейного внимания не находятся вместе внутри одного оператора softmax, мы можем использовать ассоциативность матричного умножения и считать произведение , сначала умножив на . Это и даёт желанную линейную вычислительную сложность.

Однако у такого подхода есть существенные недостатки, из-за которых модели класса линейных трансформеров не получили широкого распространения в реальных проектах. Во-первых, несмотря на асимптотическую линейность, почти все такие модели работают гораздо медленнее стандартного Трансформера в большинстве сценариев, и преимущество появляется только на очень длинных последовательностях. Это происходит потому, что для имитации поведения функции softmax используются очень вычислительно-неэффективные и затратные преобразования , такие как совокупность случайных проекций (1, 2), дополнительные MLP-слои или разложение строк и в многомерный ряд Тейлора. Во-вторых, линейные трансформеры не позволяют достичь такого же качества моделирования, как стандартный.

State Space Models

Вторая волна архитектур (с 2023 по настоящее время) представляет собой решения на основе давно известных в науке об обработке сигналов моделях пространства состояний — State Space Models (SSMs), адаптированных для нейронных сетей в замечательной серии работ про S4. Первые модели этого класса хорошо показывали себя в обработке длинных последовательностей, например звуковых сигналов или задач из бенчмарка LRA, о котором мы ещё поговорим ниже. Однако масштабные эксперименты по языковому моделированию c архитектурами серии S4 не проводились. В последующих работах некоторые параметры в SSM сделали не статичными, а зависящими от входных данных, что позволило обеспечить селективность запоминания информации, накопленной к моменту $t$ из предыдущих токенов последовательности. Этот механизм по своему действию похож на работу softmax, и новые модели, работающие по этому принципу, такие как завирусившаяся Mamba (также обзор на русском) или менее известные в широких кругах Gated Linear Attention и xLSTM, оказались конкурентными с Трансформером в авторегрессионном языковом моделировании и других задачах. 

Основной механизм в таких архитектурах разные авторы называют как хотят: и как SSMs, и как Linear RNNs (линейные рекуррентные нейронные сети), и даже как линейное внимание (хотя, с моей точки зрения, это не совсем правильно), но фундаментально это один и тот же алгоритм, который можно описать формулой:

где  и — проекции эмбеддинга токена , как и у Трансформера, а — та самая матрица «пространства состояний» из расшифровки аббревиатуры SSM. Параметр — это ключевой компонент архитектуры, который называется у разных авторов gating или decay (к сожалению, я не могу придумать красивый и подходящий по смыслу перевод этих понятий на русский язык). Он может принимать скалярное или матричное значение, но обязательно со спектральной нормой (максимальным по модулю собственным значением) меньше 1, что обеспечивает численную стабильность и фокусировку на недалёком прошлом в отсутствие softmax или матрицы нормировки . Также в Мамбе и других новых моделях этот параметр зависит от , что и обеспечивает упомянутую выше селективность.

Если убрать из формулы SSM параметр decay/gating, учесть нормировку и добавить функцию , то получится формула рекуррентного представления линейного внимания, поэтому многие считают такие архитектуры современной версией линейных трансформеров. Однако, как мы увидели, механизмы стабилизации у этих двух классов моделей разные, и у линейных трансформеров также нет математического аппарата для селективности и сосредоточения на недалёком прошлом. Тем не менее, оба алгоритма могут работать в рекуррентной форме с линейной вычислительной сложностью.

Но эта исключительно рекуррентная природа SSM является их недостатком, так как из коробки они могут быть применены только для авторегрессионного моделирования. Стандартный и линейные трансформеры могут использоваться как для авторегрессионного, так и двунаправленного (bidirectional) моделирования. Двунаправленный подход является выигрышным для любых задач, в которых вся последовательность (например, текст или изображение) известна и нужно сделать какой-то вывод о ней, к примеру, классифицировать документ или нарисовать bounding box предмета на картинке. Одни из самых известных и популярных транформерных архитектур, таких как BERT и ViT, являются двунаправленными. 

Вторым существенным недостатком селективных SSM является то, что для возможности их обучения в эффективном параллельном режиме с линейной вычислительной сложностью требуется писать низкоуровневый код, специфический для конкретного оборудования. На практике это означает, что авторы пишут реализации своих моделей на низкоуровневом языке CUDA или чуть более абстрактном фреймворке Triton, и они работают только на последних поколениях видеокарт NVIDIA. Напротив, softmax-внимание Трансформера с этой точки зрения — более простой алгоритм, который можно реализовать в несколько десятков строк кода и на чистом PyTorch, даже на NumPy, и он будет работать со скоростью одного порядка как при обучении, так и на инференсе, на любых платформах, в том числе любых процессорах и видеокартах AMD.

DenseAttention: убираем недостатки обоих подходов

В своей предыдущей работе (1, 2) я представил новую архитектуру DenseAttention Network (DANet), которая одновременно вычислительно эффективна в обоих сценариях (обучении и инференсе), работает на любом оборудовании, поддерживаемом PyTorch, и превосходит по качеству Трансформер и многие другие модели в языковом моделировании (таблица с результатами на бенчмарке GLUE) и других задачах (таблица с результатами на LRA).

При этом механизм DenseAttention является возведённым в предел упрощением внимания Транформера:

Здесь нет ни softmax, как в стандартном Трансформере, ни сложных ядерных функций и матрицы нормирования , как в линейных трансформерах, ни коэффициентов gating/decay, как в линейных SSM и RNN, — есть только чистые матричные умножения. Благодаря этому DANet не только имеет теоретическую линейную вычислительную сложность, но и быстрее на практике (таблица 6 из статьи), чем оптимизированные низкоуровневые реализации softmax-внимания даже при малой длине входной последовательности. При этом алгоритм реализован на чистом PyTorch без включений Cuda-кода. 

Полное описание архитектуры сети и причин, почему при отсутствии любых стабилизирующих элементов в механизме внимания она не превращает входные элементы в и NaNы, требует отдельного поста, а пока предлагаю почитать саму статью (1, 2). Для нашего текущего исследования представляет интерес другой вопрос: почему самый простой механизм, который только можно себе представить, даёт лучшее качество, чем современные модификации Трансформера? Причём в реальных условиях: DANet-BERT имеет 335 миллионов параметров, как и одни из самых популярных до сих пор в индустрии BERT и ROBERTa, и обучался на 500 миллиардах токенов.

Ответ на этот вопрос может привести к раскрытию скрытого потенциала в скорости и качестве у широкого класса существующих архитектур (может даже и у самого Трансформера), и к созданию новых архитектур для применения в различных специализированных задачах. 

Выделяем волшебную формулу

Мы с командой проанализировали основные элементы и дополнительные расширения архитектуры DANet, и выделили несколько потенциальных «виновников» прироста качества:

1. уменьшение количества голов внимания;

2. использование чередования слоёв локального и глобального внимания;

3. включение современных вариантов относительных позиционных эмбеддингов (Relative Positional Embeddings — RelPE).

Рассмотрим их по порядку.

Количество голов В стандартных трансформерных BERT и GPT c 335 миллионами параметров используется 16 голов внимания с размерностью 64. При таком же количестве параметров, модели DANet в экспериментах по языковому моделированию имеют только одну голову с размерностью 1024. Абляционное исследование (таблица 3 из статьи) показывает, что качество моделирования значительно ухудшается при увеличении количества голов, а одна голова большой размерности является оптимальным вариантом. На самом деле, такая ситуация — это двойная победа, потому что увеличение размерности матриц приводит к лучшей вычислительной интенсивности и большему потреблению оборудования. 

Сейчас мы проверяем, приводит ли уменьшение количества голов до одной к улучшению качества для DANet и других архитектур на задачах из бенчмарка Long Range Arena (LRA), который проверяет способности архитектур к моделированию длинных последовательностей на задачах естественного языка (NLP), компьютерного зрения (CV) и абстрактной логики. Предварительные результаты, полученные на DANet, интригующие: единственная голова гораздо лучше. Но это верно при условии предобучения модели на задаче MLM (Masked Language Modeling), похожем на обучение BERT, на тех же датасетах, которые используются и для основных задач бенчмарка LRA. При отсутствии предобучения стандартное количество голов, определённое для этих задач, показывает себя лучше одной. Но этот результат всё равно хуже, чем самый лучший результат, который получается именно при комбинации предобучение + одна голова.

Локальное и глобальное внимание. В статье «Matmuls are Enough» я также ввожу дополнение к архитектуре DANet: использование простой и вычислительно эффективной схемы слоёв нейронной сети с чередованием двух слоёв локального и одного слоя глобального внимания (см. рисунок). Слой Local просто разбивает последовательность на непересекающиеся куски длиной и считает DenseAttention внутри них. Следующий слой LocalShifted просто сдвигает границы подпоследовательностей на относительно предыдущего и, в остальном, считается так же. Последний слой Global — это просто обычный DenseAttention по всей последовательности. Похожие механизмы локального внимания известны давно (см., например, LongFormer и Block-Sparse FlashAttention), и они успешно применяются в современных трансформерных моделях, таких как Mistral и Gemma. В этих моделях глобальное внимание в части или всех слоях заменено на локальное, чтобы снизить временные затраты на вычисления, ведь сложность локального внимания — линейная. Однако включение локального внимания в модели с и без того линейной временной сложностью не приводит к выигрышу в скорости.

Причина, по которой я использую эту схему в своей модели, другая — в архитектурах с линейной вычислительной сложностью из-за отсутствия softmax нет селективного механизма. Это значит, что коэффициент внимания для всех токенов последовательности распределён более равномерно, чем в стандартном Трансформере, где он принимает значения гораздо ближе к 0 или 1 (более подробно это описано, например, здесь). Ограничение на длину окна в части слоёв позволяет алгоритму лучше концентрироваться на токенах в окрестности того, для которого считается внимание.

Использование схемы Local-LocalShifted-Global приводит к значительному улучшению качества моделирования DANet во всех рассмотренных задачах (таблица 9 из статьи). Однако эту схему можно применить не только к данной архитектуре обработки последовательностей, но и вообще к любой. Сейчас мы тестируем её на других моделях и тоже видим прирост качества.

RoPE и RelPE. Rotary Embeddings (RoPE) в последние годы доказали своё преимущество перед другими типами позиционных эмбеддингов и стали неотъемлемой частью современных трансформерных архитектур. Я даже не могу вспомнить ни одной языковой модели 2024 или 2025 года, которая бы не использовала их. Проводя эксперименты с архитектурой DANet, я тоже пришёл к тому, что интеграция относительных эмбеддингов — RoPE или другого, упрощённого варианта Cosine RelPE, предложенного мной в статье для большей вычислительной эффективности, — даёт существенный прирост в метриках качества за фиксированное количество шагов обучения или, что эквивалентно, обеспечивает более быструю сходимость. Однако простые линейные трансформеры до сих пор не пробовали объединить с RoPe. 

Заключение

Мы надеемся доказать, что с применением этих и других трюков простые линейные механизмы внимания, вроде описанного в статье «Transformers Are RNNs», становятся конкурентными с Трансформером и SSMs. Можно даже получить SOTA (state-of-the-art) результат на каких-нибудь бенчмарках. Это важно как с практической точки зрения (лёгкость получения алгоритмов, обеспечивающих высокое качество моделирования на длинных последовательностях за линейное время), так и теоретической (возможно, у нас была достойная замена Трансформеру ещё 5 лет назад, но мы применяли к ней те же правила — глобальный контекст, большое количество голов и т.д., — которые оказались неподходящими). Так что stay tuned!

Мы разрабатываем фреймворк для обучения и исследований эффективных альтернатив Трансформеру в реальном времени. Проект, а также код DANet доступны здесь 

Если у вас есть вопросы или предложения, пишите мне в тг @andrewargatkiny, буду рад ответить.